diff --git a/pipelines/Tr006_V2/README.md b/pipelines/Tr006_V2/README.md new file mode 100644 index 0000000000000000000000000000000000000000..45db8e5e7eca60d4db341d0b3252822b93f971aa --- /dev/null +++ b/pipelines/Tr006_V2/README.md @@ -0,0 +1,129 @@ +# Tr006 V2 Segment pipeline bundle + +这个目录对应 Hugging Face 仓库 `Gbone3176/interval` 中的 +`pipelines/Tr006_V2/`。权重位于同仓库的 +`weights/Tr006_V2/checkpoint-6013/`。 + +## 范围 + +包含 Tr006 V2 实际使用的完整处理链： + +1. 原始 HeiCo 视频生成冻结蓝屏区间。 +2. 官方原始 Segment parquet 和完整视频生成 V1 fixed80 数据。 +3. 使用冻结蓝屏区间生成 V2 blue-clean fixed80 数据及 ms-swift `train.jsonl`。 +4. ms-swift 绝对时间戳模板、训练入口和实际 Tr006 配置。 +5. checkpoint-6013 LoRA merge。 +6. Segment Docker 的预处理、推理及答案后处理。 +7. 本地 vLLM 评测和时间容差复算脚本，仅用于本地核对；它们不是官方 evaluator。 + +不包含官方原始视频、官方 parquet、base model 或合并后的模型。Farm04 必须使用与原服务器相同版本的官方原始数据和同一个 Qwen3.5-9B base model。 + +## 原始数据目录契约 + +```text +RAW_ROOT/ +├── heico-focus-vqa/ +│ ├── data/segment/{train,test}.parquet +│ └── videos/* +└── lapchole-focus-vqa/ + ├── data/segment/{train,test}.parquet + └── videos/* +``` + +先用 `manifest/pipeline_manifest.json` 中的 parquet SHA256 核对版本。 + +## 从完全原始数据生成训练输入 + +### 1. 可重复生成蓝屏区间 + +```bash +python project/data_preprocessing/process_blue_screen_dataset.py \ + --input-dir "$RAW_ROOT/heico-focus-vqa/videos" \ + --output-dir "$WORK_ROOT/blue_screen_detected" \ + --sample-seconds 1 \ + --min-static-seconds 6 \ + --blue-ratio 0.45 \ + --pad-seconds 5 \ + --coarse-size 96 \ + --fine-size 160 +``` + +Tr006 的精确复现应使用 bundle 内冻结的 +`annotations/blue_screen/data/intervals.csv`。重新检测用于验证原始到标注的完整链路；只有生成文件的 SHA256 与 manifest 一致时，才可替代冻结文件。 + +### 2. 原始数据生成 V1 fixed80 + +```bash +python project/scripts/preprocess_segmentqa_fixed_frames.py \ + --data-root "$RAW_ROOT" \ + --output-dir "$WORK_ROOT/orena_segmentqa_fixed80_420x224_absmeta_v1/clean" \ + --splits train test \ + --frames 80 \ + --width 420 \ + --height 224 \ + --workers 8 +``` + +### 3. V1 生成 Tr006 V2 训练数据 + +```bash +python project/scripts/preprocess_segmentqa_blueclean_v2.py \ + --data-root "$RAW_ROOT" \ + --v1-dir "$WORK_ROOT/orena_segmentqa_fixed80_420x224_absmeta_v1/clean" \ + --blue-intervals "$BUNDLE_ROOT/annotations/blue_screen/data/intervals.csv" \ + --output-dir "$WORK_ROOT/orena_segmentqa_fixed80_420x224_absmeta_v2/clean" \ + --splits train test \ + --frames 80 \ + --width 420 \ + --height 224 \ + --workers 8 \ + --min-train-valid-frames 16 +``` + +验收值：18,478 个 segment、1,215 个蓝屏重叠 segment、13,746 条训练问题、6,254 条测试问题、80 帧、420x224、绝对毫秒 PTS。最终训练入口是 V2 目录中的 `clean/train.jsonl`。 + +`project/` 保留原项目相对布局；其中已经包含 Tr006 实际使用的 +`ms-swift/swift/template/templates/qwen.py`。如果改用另一个 ms-swift checkout，必须把该文件覆盖到 commit +`b66edaeba3c484c431d785f95a7fc81cffa530c0` 的对应位置，否则预计算绝对时间戳不会进入 Qwen3VL processor。 + +## Merge checkpoint-6013 + +```bash +python project/scripts/merge_lora_model.py \ + --base "$BASE_MODEL" \ + --adapter "$HF_ROOT/weights/Tr006_V2/checkpoint-6013" \ + --output "$OFFICIAL_SEGMENT_DIR/resources/model" \ + --max-shard-size 4GB +``` + +验收 `resources/model/merge_report.json`： + +- `remaining_lora_modules_after_merge == 0` +- language model、visual、visual merger 的代表权重 delta 均非零 +- 输出含完整 processor 配置和 safetensors index + +## 使用官方目录重新打包 + +继续使用官方原始 `do_build.sh`、`do_test_run.sh` 和 `do_save.sh`。把 bundle 中以下三个文件覆盖到官方 Segment 目录： + +```bash +cp docker/segment-algorithm/inference.py "$OFFICIAL_SEGMENT_DIR/inference.py" +cp docker/segment-algorithm/requirements.txt "$OFFICIAL_SEGMENT_DIR/requirements.txt" +cp docker/segment-algorithm/Dockerfile "$OFFICIAL_SEGMENT_DIR/Dockerfile" +``` + +完成 merge 后执行： + +```bash +cd "$OFFICIAL_SEGMENT_DIR" +bash do_build.sh +bash do_test_run.sh +bash do_save.sh +``` + +不要使用当前线上版本的 `np.linspace` 推理预处理作为“已与 V2 对齐”的证据。必须先在公开样本上比较训练 V2 的采样 PTS、420x224 帧、processor 输入和 Docker 输入；`docker/segment-algorithm/inference.py` 是已提交版本的事实快照，保留它是为了复现和定位差异。 + +## 本地评测边界 + +`project/scripts/eval_segmentqa_vllm.py` 和 +`project/scripts/evaluate_segmentqa_predictions.py` 使用本地 `correct()` 和时间容差代理，不是官方十桶 ID/OOD evaluator，也没有运行官方 LLM judge。禁止把其分数直接当成 leaderboard score。 diff --git a/pipelines/Tr006_V2/annotations/blue_screen/data/intervals.csv b/pipelines/Tr006_V2/annotations/blue_screen/data/intervals.csv new file mode 100644 index 0000000000000000000000000000000000000000..8464236ff0f082c56706986c2fd322bdde3115f2 --- /dev/null +++ b/pipelines/Tr006_V2/annotations/blue_screen/data/intervals.csv @@ -0,0 +1,324 @@ +video_id,video_filename,interval_index,start_frame,end_frame,start_seconds,end_seconds,duration_seconds,start_time,end_time,frame_count,mean_blue_ratio,confidence,exact_repeated_hash +0000,0000 - Heico - Prokto - 1.avi,1,5524,5938,220.96,237.56,16.6,00:03:40.960,00:03:57.560,415,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0000,0000 - Heico - Prokto - 1.avi,2,9296,9548,371.84,381.96,10.12,00:06:11.840,00:06:21.960,253,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0000,0000 - Heico - Prokto - 1.avi,3,10249,10494,409.96,419.8,9.84,00:06:49.960,00:06:59.800,246,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0000,0000 - Heico - Prokto - 1.avi,4,105966,106309,4238.64,4252.4,13.76,01:10:38.640,01:10:52.400,344,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0000,0000 - Heico - Prokto - 1.avi,5,153266,153973,6130.64,6158.96,28.32,01:42:10.640,01:42:38.960,708,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0000,0000 - Heico - Prokto - 1.avi,6,216719,217948,8668.76,8717.96,49.2,02:24:28.760,02:25:17.960,1230,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0000,0000 - Heico - Prokto - 1.avi,7,232574,311655,9302.96,12466.24,3163.28,02:35:02.960,03:27:46.240,79082,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0001,0001 - Heico - Prokto - 2.avi,1,110844,111882,4433.76,4475.32,41.56,01:13:53.760,01:14:35.320,1039,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0001,0001 - Heico - Prokto - 2.avi,2,112792,113563,4511.68,4542.56,30.88,01:15:11.680,01:15:42.560,772,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0001,0001 - Heico - Prokto - 2.avi,3,114462,115183,4578.48,4607.36,28.88,01:16:18.480,01:16:47.360,722,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0001,0001 - Heico - Prokto - 2.avi,4,137203,138843,5488.12,5553.76,65.64,01:31:28.120,01:32:33.760,1641,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0001,0001 - Heico - Prokto - 2.avi,5,147830,148183,5913.2,5927.36,14.16,01:38:33.200,01:38:47.360,354,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0001,0001 - Heico - Prokto - 2.avi,6,154539,155178,6181.56,6207.16,25.6,01:43:01.560,01:43:27.160,640,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0001,0001 - Heico - Prokto - 2.avi,7,222864,223263,8914.56,8930.56,16.0,02:28:34.560,02:28:50.560,400,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0001,0001 - Heico - Prokto - 2.avi,8,237013,316761,9480.52,12670.48,3189.96,02:38:00.520,03:31:10.480,79749,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,1,5009,5212,200.36,208.52,8.16,00:03:20.360,00:03:28.520,204,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,2,5326,5704,213.04,228.2,15.16,00:03:33.040,00:03:48.200,379,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,3,26052,27664,1042.08,1106.6,64.52,00:17:22.080,00:18:26.600,1613,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,4,115654,115924,4626.16,4637.0,10.84,01:17:06.160,01:17:17.000,271,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,5,136596,136848,5463.84,5473.96,10.12,01:31:03.840,01:31:13.960,253,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,6,144722,144905,5788.88,5796.24,7.36,01:36:28.880,01:36:36.240,184,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,7,159926,217858,6397.04,8714.36,2317.32,01:46:37.040,02:25:14.360,57933,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,8,227798,229000,9111.92,9160.04,48.12,02:31:51.920,02:32:40.040,1203,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,9,248690,251825,9947.6,10073.04,125.44,02:45:47.600,02:47:53.040,3136,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,10,252911,254146,10116.44,10165.88,49.44,02:48:36.440,02:49:25.880,1236,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0002,0002 - Heico - Prokto - 3.avi,11,255848,256281,10233.92,10251.28,17.36,02:50:33.920,02:50:51.280,434,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,1,643,869,25.72,34.8,9.08,00:00:25.720,00:00:34.800,227,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,2,1267,1562,50.68,62.52,11.84,00:00:50.680,00:01:02.520,296,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,3,11885,12122,475.4,484.92,9.52,00:07:55.400,00:08:04.920,238,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,4,82913,83622,3316.52,3344.92,28.4,00:55:16.520,00:55:44.920,710,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,5,127912,128071,5116.48,5122.88,6.4,01:25:16.480,01:25:22.880,160,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,6,141561,200111,5662.44,8004.48,2342.04,01:34:22.440,02:13:24.480,58551,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,7,205346,206231,8213.84,8249.28,35.44,02:16:53.840,02:17:29.280,886,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,8,233424,237433,9336.96,9497.36,160.4,02:35:36.960,02:38:17.360,4010,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,9,254178,255182,10167.12,10207.32,40.2,02:49:27.120,02:50:07.320,1005,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0003,0003 - Heico - Prokto - 4.avi,10,255278,255922,10211.12,10236.92,25.8,02:50:11.120,02:50:36.920,645,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,1,811,1005,32.44,40.24,7.8,00:00:32.440,00:00:40.240,195,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,2,2507,2827,100.28,113.12,12.84,00:01:40.280,00:01:53.120,321,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,3,2911,3179,116.44,127.2,10.76,00:01:56.440,00:02:07.200,269,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,4,5065,5270,202.6,210.84,8.24,00:03:22.600,00:03:30.840,206,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,5,5487,5670,219.48,226.84,7.36,00:03:39.480,00:03:46.840,184,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,6,153423,153582,6136.92,6143.32,6.4,01:42:16.920,01:42:23.320,160,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,7,166711,237718,6668.44,9508.76,2840.32,01:51:08.440,02:38:28.760,71008,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,8,274728,274887,10989.12,10995.52,6.4,03:03:09.120,03:03:15.520,160,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,9,279748,280329,11189.92,11213.2,23.28,03:06:29.920,03:06:53.200,582,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,10,281024,281672,11240.96,11266.92,25.96,03:07:20.960,03:07:46.920,649,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,11,295193,298502,11807.72,11940.12,132.4,03:16:47.720,03:19:00.120,3310,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,12,299039,299984,11961.56,11999.4,37.84,03:19:21.560,03:19:59.400,946,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,13,301545,302528,12061.8,12101.16,39.36,03:21:01.800,03:21:41.160,984,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,14,304448,305829,12177.92,12233.2,55.28,03:22:57.920,03:23:53.200,1382,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,15,306273,312447,12250.92,12497.92,247.0,03:24:10.920,03:28:17.920,6175,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,16,313573,313830,12542.92,12553.24,10.32,03:29:02.920,03:29:13.240,258,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0004,0004 - Heico - Prokto - 5.avi,17,315328,315503,12613.12,12620.16,7.04,03:30:13.120,03:30:20.160,176,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,1,63,977,2.52,39.12,36.6,00:00:02.520,00:00:39.120,915,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,2,2600,2834,104.0,113.4,9.4,00:01:44.000,00:01:53.400,235,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,3,88472,88932,3538.88,3557.32,18.44,00:58:58.880,00:59:17.320,461,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,4,89225,89747,3569.0,3589.92,20.92,00:59:29.000,00:59:49.920,523,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,5,89971,90170,3598.84,3606.84,8.0,00:59:58.840,01:00:06.840,200,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,6,117330,166097,4693.2,6643.92,1950.72,01:18:13.200,01:50:43.920,48768,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,7,168608,168844,6744.32,6753.8,9.48,01:52:24.320,01:52:33.800,237,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,8,184972,185273,7398.88,7410.96,12.08,02:03:18.880,02:03:30.960,302,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0005,0005 - Heico - Prokto - 6.avi,9,199269,199539,7970.76,7981.6,10.84,02:12:50.760,02:13:01.600,271,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,1,54370,54766,2174.8,2190.68,15.88,00:36:14.800,00:36:30.680,397,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,2,68517,68799,2740.68,2752.0,11.32,00:45:40.680,00:45:52.000,283,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,3,131005,192543,5240.2,7701.76,2461.56,01:27:20.200,02:08:21.760,61539,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,4,192597,192916,7703.88,7716.68,12.8,02:08:23.880,02:08:36.680,320,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,5,196217,197106,7848.68,7884.28,35.6,02:10:48.680,02:11:24.280,890,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,6,198352,198542,7934.08,7941.72,7.64,02:12:14.080,02:12:21.720,191,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,7,204030,204189,8161.2,8167.6,6.4,02:16:01.200,02:16:07.600,160,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,8,222489,222898,8899.56,8915.96,16.4,02:28:19.560,02:28:35.960,410,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,9,239305,239671,9572.2,9586.88,14.68,02:39:32.200,02:39:46.880,367,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,10,241199,241485,9647.96,9659.44,11.48,02:40:47.960,02:40:59.440,287,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0006,0006 - Heico - Prokto - 7.avi,11,247032,252252,9881.28,10090.12,208.84,02:44:41.280,02:48:10.120,5221,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,1,2119,2344,84.76,93.8,9.04,00:01:24.760,00:01:33.800,226,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,2,2520,2930,100.8,117.24,16.44,00:01:40.800,00:01:57.240,411,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,3,77279,77433,3091.16,3097.36,6.2,00:51:31.160,00:51:37.360,155,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,4,125114,126789,5004.56,5071.6,67.04,01:23:24.560,01:24:31.600,1676,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,5,150927,209117,6037.08,8364.72,2327.64,01:40:37.080,02:19:24.720,58191,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,6,209372,209595,8374.88,8383.84,8.96,02:19:34.880,02:19:43.840,224,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,7,237291,237838,9491.64,9513.56,21.92,02:38:11.640,02:38:33.560,548,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,8,238028,238568,9521.12,9542.76,21.64,02:38:41.120,02:39:02.760,541,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,9,238736,238982,9549.44,9559.32,9.88,02:39:09.440,02:39:19.320,247,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0007,0007 - Heico - Prokto - 8.avi,10,252643,254287,10105.72,10171.52,65.8,02:48:25.720,02:49:31.520,1645,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,1,2459,2677,98.36,107.12,8.76,00:01:38.360,00:01:47.120,219,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,2,2871,3130,114.84,125.24,10.4,00:01:54.840,00:02:05.240,260,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,3,5009,5173,200.36,206.96,6.6,00:03:20.360,00:03:26.960,165,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,4,5444,5697,217.76,227.92,10.16,00:03:37.760,00:03:47.920,254,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,5,166621,237773,6664.84,9510.96,2846.12,01:51:04.840,02:38:30.960,71153,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,6,238250,238399,9530.0,9536.0,6.0,02:38:50.000,02:38:56.000,150,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,7,279716,280051,11188.64,11202.08,13.44,03:06:28.640,03:06:42.080,336,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,8,280996,281540,11239.84,11261.64,21.8,03:07:19.840,03:07:41.640,545,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,9,295157,298333,11806.28,11933.36,127.08,03:16:46.280,03:18:53.360,3177,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,10,299481,299873,11979.24,11994.96,15.72,03:19:39.240,03:19:54.960,393,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,11,301505,302399,12060.2,12096.0,35.8,03:21:00.200,03:21:36.000,895,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,12,304394,305695,12175.76,12227.84,52.08,03:22:55.760,03:23:47.840,1302,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,13,306196,310438,12247.84,12417.56,169.72,03:24:07.840,03:26:57.560,4243,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,14,313476,313750,12539.04,12550.04,11.0,03:28:59.040,03:29:10.040,275,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0008,0008 - Heico - Prokto - 9.avi,15,315281,315430,12611.24,12617.24,6.0,03:30:11.240,03:30:17.240,150,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,1,7060,8007,282.4,320.32,37.92,00:04:42.400,00:05:20.320,948,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,2,8244,8518,329.76,340.76,11.0,00:05:29.760,00:05:40.760,275,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,3,10905,11723,436.2,468.96,32.76,00:07:16.200,00:07:48.960,819,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,4,12679,12928,507.16,517.16,10.0,00:08:27.160,00:08:37.160,250,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,5,13404,13908,536.16,556.36,20.2,00:08:56.160,00:09:16.360,505,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,6,53068,58670,2122.72,2346.84,224.12,00:35:22.720,00:39:06.840,5603,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,7,63976,65220,2559.04,2608.84,49.8,00:42:39.040,00:43:28.840,1245,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,8,65405,65853,2616.2,2634.16,17.96,00:43:36.200,00:43:54.160,449,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,9,66390,67030,2655.6,2681.24,25.64,00:44:15.600,00:44:41.240,641,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,10,85249,86243,3409.96,3449.76,39.8,00:56:49.960,00:57:29.760,995,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,11,91754,92359,3670.16,3694.4,24.24,01:01:10.160,01:01:34.400,606,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,12,92636,93048,3705.44,3721.96,16.52,01:01:45.440,01:02:01.960,413,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,13,94012,94245,3760.48,3769.84,9.36,01:02:40.480,01:02:49.840,234,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,14,125101,126017,5004.04,5040.72,36.68,01:23:24.040,01:24:00.720,917,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,15,126104,126615,5044.16,5064.64,20.48,01:24:04.160,01:24:24.640,512,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,16,145892,146607,5835.68,5864.32,28.64,01:37:15.680,01:37:44.320,716,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,17,147498,148638,5899.92,5945.56,45.64,01:38:19.920,01:39:05.560,1141,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,18,151144,153672,6045.76,6146.92,101.16,01:40:45.760,01:42:26.920,2529,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,19,153819,154897,6152.76,6195.92,43.16,01:42:32.760,01:43:15.920,1079,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,20,156864,157369,6274.56,6294.8,20.24,01:44:34.560,01:44:54.800,506,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,21,158007,158420,6320.28,6336.84,16.56,01:45:20.280,01:45:36.840,414,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,22,171517,171713,6860.68,6868.56,7.88,01:54:20.680,01:54:28.560,197,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,23,185776,187204,7431.04,7488.2,57.16,02:03:51.040,02:04:48.200,1429,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,24,187469,187733,7498.76,7509.36,10.6,02:04:58.760,02:05:09.360,265,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,25,206430,207747,8257.2,8309.92,52.72,02:17:37.200,02:18:29.920,1318,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,26,207944,208182,8317.76,8327.32,9.56,02:18:37.760,02:18:47.320,239,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,27,214536,214878,8581.44,8595.16,13.72,02:23:01.440,02:23:15.160,343,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,28,214972,215438,8598.88,8617.56,18.68,02:23:18.880,02:23:37.560,467,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,29,216309,216908,8652.36,8676.36,24.0,02:24:12.360,02:24:36.360,600,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,30,219259,219518,8770.36,8780.76,10.4,02:26:10.360,02:26:20.760,260,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,31,223627,223803,8945.08,8952.16,7.08,02:29:05.080,02:29:12.160,177,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,32,229414,229663,9176.56,9186.56,10.0,02:32:56.560,02:33:06.560,250,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,33,280676,283232,11227.04,11329.32,102.28,03:07:07.040,03:08:49.320,2557,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,34,308970,373914,12358.8,14956.6,2597.8,03:25:58.800,04:09:16.600,64945,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,35,374164,374543,14966.56,14981.76,15.2,04:09:26.560,04:09:41.760,380,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,36,387421,388107,15496.84,15524.32,27.48,04:18:16.840,04:18:44.320,687,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,37,404192,404443,16167.68,16177.76,10.08,04:29:27.680,04:29:37.760,252,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,38,407105,407310,16284.2,16292.44,8.24,04:31:24.200,04:31:32.440,206,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,39,432003,441298,17280.12,17651.96,371.84,04:48:00.120,04:54:11.960,9296,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0009,0009 - Heico - Prokto - 10.avi,40,442866,443022,17714.64,17720.92,6.28,04:55:14.640,04:55:20.920,157,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,1,41235,41710,1649.4,1668.44,19.04,00:27:29.400,00:27:48.440,476,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,2,49196,50101,1967.84,2004.08,36.24,00:32:47.840,00:33:24.080,906,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,3,157223,158078,6288.92,6323.16,34.24,01:44:48.920,01:45:23.160,856,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,4,160177,161891,6407.08,6475.68,68.6,01:46:47.080,01:47:55.680,1715,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,5,165333,165802,6613.32,6632.12,18.8,01:50:13.320,01:50:32.120,470,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,6,265363,265733,10614.52,10629.36,14.84,02:56:54.520,02:57:09.360,371,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,7,265907,336405,10636.28,13456.24,2819.96,02:57:16.280,03:44:16.240,70499,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,8,354902,355116,14196.08,14204.68,8.6,03:56:36.080,03:56:44.680,215,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,9,371415,371767,14856.6,14870.72,14.12,04:07:36.600,04:07:50.720,353,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,10,373541,373981,14941.64,14959.28,17.64,04:09:01.640,04:09:19.280,441,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,11,376427,376726,15057.08,15069.08,12.0,04:10:57.080,04:11:09.080,300,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,12,397827,398219,15913.08,15928.8,15.72,04:25:13.080,04:25:28.800,393,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0010,0010 - Heico - Rektum - 1.avi,13,421164,429147,16846.56,17165.92,319.36,04:40:46.560,04:46:05.920,7984,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,1,7043,7277,281.72,291.12,9.4,00:04:41.720,00:04:51.120,235,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,2,176925,177260,7077.0,7090.44,13.44,01:57:57.000,01:58:10.440,336,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,3,183595,232058,7343.8,9282.36,1938.56,02:02:23.800,02:34:42.360,48464,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,4,239127,240057,9565.08,9602.32,37.24,02:39:25.080,02:40:02.320,931,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,5,240627,241173,9625.08,9646.96,21.88,02:40:25.080,02:40:46.960,547,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,6,247248,247415,9889.92,9896.64,6.72,02:44:49.920,02:44:56.640,168,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,7,251653,254734,10066.12,10189.4,123.28,02:47:46.120,02:49:49.400,3082,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,8,256223,256422,10248.92,10256.92,8.0,02:50:48.920,02:50:56.920,200,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,9,266223,266459,10648.92,10658.4,9.48,02:57:28.920,02:57:38.400,237,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,10,309058,309278,12362.32,12371.16,8.84,03:26:02.320,03:26:11.160,221,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0011,0011 - Heico - Rektum - 2.avi,11,330794,331069,13231.76,13242.8,11.04,03:40:31.760,03:40:42.800,276,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0012,0012 - Heico - Rektum - 3.avi,1,18110,18492,724.4,739.72,15.32,00:12:04.400,00:12:19.720,383,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0012,0012 - Heico - Rektum - 3.avi,2,150983,191811,6039.32,7672.48,1633.16,01:40:39.320,02:07:52.480,40829,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0012,0012 - Heico - Rektum - 3.avi,3,211660,211895,8466.4,8475.84,9.44,02:21:06.400,02:21:15.840,236,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0013,0013 - Heico - Rektum - 4.avi,1,8189,8713,327.56,348.56,21.0,00:05:27.560,00:05:48.560,525,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0013,0013 - Heico - Rektum - 4.avi,2,9174,9696,366.96,387.88,20.92,00:06:06.960,00:06:27.880,523,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0013,0013 - Heico - Rektum - 4.avi,3,12253,13148,490.12,525.96,35.84,00:08:10.120,00:08:45.960,896,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0013,0013 - Heico - Rektum - 4.avi,4,16396,16578,655.84,663.16,7.32,00:10:55.840,00:11:03.160,183,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0013,0013 - Heico - Rektum - 4.avi,5,164401,164911,6576.04,6596.48,20.44,01:49:36.040,01:49:56.480,511,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0013,0013 - Heico - Rektum - 4.avi,6,168791,208270,6751.64,8330.84,1579.2,01:52:31.640,02:18:50.840,39480,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,1,1061,1842,42.44,73.72,31.28,00:00:42.440,00:01:13.720,782,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,2,7887,8186,315.48,327.48,12.0,00:05:15.480,00:05:27.480,300,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,3,9287,9578,371.48,383.16,11.68,00:06:11.480,00:06:23.160,292,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,4,12362,12576,494.48,503.08,8.6,00:08:14.480,00:08:23.080,215,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,5,12932,13344,517.28,533.8,16.52,00:08:37.280,00:08:53.800,413,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,6,14932,15101,597.28,604.08,6.8,00:09:57.280,00:10:04.080,170,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,7,20627,20781,825.08,831.28,6.2,00:13:45.080,00:13:51.280,155,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,8,46852,47131,1874.08,1885.28,11.2,00:31:14.080,00:31:25.280,280,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,9,83957,84631,3358.28,3385.28,27.0,00:55:58.280,00:56:25.280,675,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,10,152532,152808,6101.28,6112.36,11.08,01:41:41.280,01:41:52.360,277,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,11,159772,159947,6390.88,6397.92,7.04,01:46:30.880,01:46:37.920,176,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,12,161457,161771,6458.28,6470.88,12.6,01:47:38.280,01:47:50.880,315,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,13,164407,164931,6576.28,6597.28,21.0,01:49:36.280,01:49:57.280,525,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,14,165100,165406,6604.0,6616.28,12.28,01:50:04.000,01:50:16.280,307,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,15,170881,171031,6835.24,6841.28,6.04,01:53:55.240,01:54:01.280,151,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,16,214275,256687,8571.0,10267.52,1696.52,02:22:51.000,02:51:07.520,42413,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,17,269272,270063,10770.88,10802.56,31.68,02:59:30.880,03:00:02.560,792,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0014,0014 - Heico - Rektum - 5.avi,18,286533,300706,11461.32,12028.28,566.96,03:11:01.320,03:20:28.280,14174,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,1,98,871,3.92,34.88,30.96,00:00:03.920,00:00:34.880,774,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,2,1030,1304,41.2,52.2,11.0,00:00:41.200,00:00:52.200,275,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,3,4210,4569,168.4,182.8,14.4,00:02:48.400,00:03:02.800,360,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,4,24926,25082,997.04,1003.32,6.28,00:16:37.040,00:16:43.320,157,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,5,25652,26954,1026.08,1078.2,52.12,00:17:06.080,00:17:58.200,1303,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,6,71242,71463,2849.68,2858.56,8.88,00:47:29.680,00:47:38.560,222,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,7,72854,73042,2914.16,2921.72,7.56,00:48:34.160,00:48:41.720,189,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,8,118161,118453,4726.44,4738.16,11.72,01:18:46.440,01:18:58.160,293,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,9,154118,154421,6164.72,6176.88,12.16,01:42:44.720,01:42:56.880,304,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,10,225395,225584,9015.8,9023.4,7.6,02:30:15.800,02:30:23.400,190,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,11,238226,238394,9529.04,9535.8,6.76,02:38:49.040,02:38:55.800,169,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,12,239171,239361,9566.84,9574.48,7.64,02:39:26.840,02:39:34.480,191,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,13,252669,252901,10106.76,10116.08,9.32,02:48:26.760,02:48:36.080,233,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,14,254945,369078,10197.8,14763.16,4565.36,02:49:57.800,04:06:03.160,114134,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0015,0015 - Heico - Rektum - 6.avi,15,380814,381016,15232.56,15240.68,8.12,04:13:52.560,04:14:00.680,203,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,1,286,504,11.44,20.2,8.76,00:00:11.440,00:00:20.200,219,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,2,1535,2679,61.4,107.2,45.8,00:01:01.400,00:01:47.200,1145,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,3,3366,3528,134.64,141.16,6.52,00:02:14.640,00:02:21.160,163,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,4,76469,77327,3058.76,3093.12,34.36,00:50:58.760,00:51:33.120,859,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,5,111617,111991,4464.68,4479.68,15.0,01:14:24.680,01:14:39.680,375,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,6,176683,226219,7067.32,9048.8,1981.48,01:57:47.320,02:30:48.800,49537,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,7,230857,231447,9234.28,9257.92,23.64,02:33:54.280,02:34:17.920,591,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,8,245796,246039,9831.84,9841.6,9.76,02:43:51.840,02:44:01.600,244,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,9,256669,257081,10266.76,10283.28,16.52,02:51:06.760,02:51:23.280,413,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,10,262393,262789,10495.72,10511.6,15.88,02:54:55.720,02:55:11.600,397,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,11,273830,274150,10953.2,10966.04,12.84,03:02:33.200,03:02:46.040,321,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,12,275856,276843,11034.24,11073.76,39.52,03:03:54.240,03:04:33.760,988,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,13,306912,317030,12276.48,12681.24,404.76,03:24:36.480,03:31:21.240,10119,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,14,325605,337482,13024.2,13499.32,475.12,03:37:04.200,03:44:59.320,11878,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,15,337785,349126,13511.4,13965.08,453.68,03:45:11.400,03:52:45.080,11342,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0016,0016 - Heico - Rektum - 7.avi,16,380708,381065,15228.32,15242.64,14.32,04:13:48.320,04:14:02.640,358,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0017,0017 - Heico - Rektum - 8.avi,1,3448,3630,137.92,145.24,7.32,00:02:17.920,00:02:25.240,183,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0017,0017 - Heico - Rektum - 8.avi,2,3746,4105,149.84,164.24,14.4,00:02:29.840,00:02:44.240,360,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0017,0017 - Heico - Rektum - 8.avi,3,5012,5189,200.48,207.6,7.12,00:03:20.480,00:03:27.600,178,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0017,0017 - Heico - Rektum - 8.avi,4,10783,11192,431.32,447.72,16.4,00:07:11.320,00:07:27.720,410,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0017,0017 - Heico - Rektum - 8.avi,5,143141,185332,5725.64,7413.32,1687.68,01:35:25.640,02:03:33.320,42192,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0017,0017 - Heico - Rektum - 8.avi,6,210088,210262,8403.52,8410.52,7.0,02:20:03.520,02:20:10.520,175,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0018,0018 - Heico - Rektum - 9.avi,1,4605,6290,184.2,251.64,67.44,00:03:04.200,00:04:11.640,1686,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0018,0018 - Heico - Rektum - 9.avi,2,196494,196665,7859.76,7866.64,6.88,02:10:59.760,02:11:06.640,172,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0018,0018 - Heico - Rektum - 9.avi,3,207338,260298,8293.52,10411.96,2118.44,02:18:13.520,02:53:31.960,52961,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0018,0018 - Heico - Rektum - 9.avi,4,267668,268079,10706.72,10723.2,16.48,02:58:26.720,02:58:43.200,412,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0018,0018 - Heico - Rektum - 9.avi,5,269857,270216,10794.28,10808.68,14.4,02:59:54.280,03:00:08.680,360,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0018,0018 - Heico - Rektum - 9.avi,6,345674,346337,13826.96,13853.52,26.56,03:50:26.960,03:50:53.520,664,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0019,0019 - Heico - Rektum - 10.avi,1,88458,88727,3538.32,3549.12,10.8,00:58:58.320,00:59:09.120,270,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0019,0019 - Heico - Rektum - 10.avi,2,101634,155868,4065.36,6234.76,2169.4,01:07:45.360,01:43:54.760,54235,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0019,0019 - Heico - Rektum - 10.avi,3,157074,157625,6282.96,6305.04,22.08,01:44:42.960,01:45:05.040,552,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0020,0020 - Heico - Sigma - 1.avi,1,76741,120650,3069.64,4826.04,1756.4,00:51:09.640,01:20:26.040,43910,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0021,0021 - Heico - Sigma - 2.avi,1,1403,1602,56.12,64.12,8.0,00:00:56.120,00:01:04.120,200,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0021,0021 - Heico - Sigma - 2.avi,2,75202,129434,3008.08,5177.4,2169.32,00:50:08.080,01:26:17.400,54233,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0021,0021 - Heico - Sigma - 2.avi,3,142160,142314,5686.4,5692.6,6.2,01:34:46.400,01:34:52.600,155,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0021,0021 - Heico - Sigma - 2.avi,4,145460,145669,5818.4,5826.8,8.4,01:36:58.400,01:37:06.800,210,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,1,4196,4772,167.84,190.92,23.08,00:02:47.840,00:03:10.920,577,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,2,58606,59151,2344.24,2366.08,21.84,00:39:04.240,00:39:26.080,546,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,3,108376,109561,4335.04,4382.48,47.44,01:12:15.040,01:13:02.480,1186,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,4,150566,151065,6022.64,6042.64,20.0,01:40:22.640,01:40:42.640,500,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,5,213564,213819,8542.56,8552.8,10.24,02:22:22.560,02:22:32.800,256,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,6,213927,214506,8557.08,8580.28,23.2,02:22:37.080,02:23:00.280,580,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,7,235981,292296,9439.24,11691.88,2252.64,02:37:19.240,03:14:51.880,56316,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,8,308248,308950,12329.92,12358.04,28.12,03:25:29.920,03:25:58.040,703,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,9,313811,314209,12552.44,12568.4,15.96,03:29:12.440,03:29:28.400,399,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0022,0022 - Heico - Sigma - 3.avi,10,331328,332053,13253.12,13282.16,29.04,03:40:53.120,03:41:22.160,726,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0023,0023 - Heico - Sigma - 4.avi,1,22393,22586,895.72,903.48,7.76,00:14:55.720,00:15:03.480,194,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0023,0023 - Heico - Sigma - 4.avi,2,23387,23797,935.48,951.92,16.44,00:15:35.480,00:15:51.920,411,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0023,0023 - Heico - Sigma - 4.avi,3,58098,58324,2323.92,2333.0,9.08,00:38:43.920,00:38:53.000,227,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0023,0023 - Heico - Sigma - 4.avi,4,89007,89196,3560.28,3567.88,7.6,00:59:20.280,00:59:27.880,190,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0023,0023 - Heico - Sigma - 4.avi,5,135563,136671,5422.52,5466.88,44.36,01:30:22.520,01:31:06.880,1109,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0023,0023 - Heico - Sigma - 4.avi,6,136822,188918,5472.88,7556.76,2083.88,01:31:12.880,02:05:56.760,52097,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0023,0023 - Heico - Sigma - 4.avi,7,189491,189698,7579.64,7587.96,8.32,02:06:19.640,02:06:27.960,208,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0023,0023 - Heico - Sigma - 4.avi,8,216348,216559,8653.92,8662.4,8.48,02:24:13.920,02:24:22.400,212,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0024,0024 - Heico - Sigma - 5.avi,1,3523,3791,140.92,151.68,10.76,00:02:20.920,00:02:31.680,269,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0024,0024 - Heico - Sigma - 5.avi,2,49240,49671,1969.6,1986.88,17.28,00:32:49.600,00:33:06.880,432,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0024,0024 - Heico - Sigma - 5.avi,3,72121,129067,2884.84,5162.72,2277.88,00:48:04.840,01:26:02.720,56947,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,1,769,971,30.76,38.88,8.12,00:00:30.760,00:00:38.880,203,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,2,5809,5986,232.36,239.48,7.12,00:03:52.360,00:03:59.480,178,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,3,53622,54258,2144.88,2170.36,25.48,00:35:44.880,00:36:10.360,637,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,4,96792,96976,3871.68,3879.08,7.4,01:04:31.680,01:04:39.080,185,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,5,119669,218770,4786.76,8750.84,3964.08,01:19:46.760,02:25:50.840,99102,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,6,228355,230160,9134.2,9206.44,72.24,02:32:14.200,02:33:26.440,1806,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,7,250996,252051,10039.84,10082.08,42.24,02:47:19.840,02:48:02.080,1056,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,8,254176,254633,10167.04,10185.36,18.32,02:49:27.040,02:49:45.360,458,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,9,294134,294341,11765.36,11773.68,8.32,03:16:05.360,03:16:13.680,208,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0025,0025 - Heico - Sigma - 6.avi,10,305932,306158,12237.28,12246.36,9.08,03:23:57.280,03:24:06.360,227,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0026,0026 - Heico - Sigma - 7.avi,1,104407,104647,4176.28,4185.92,9.64,01:09:36.280,01:09:45.920,241,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0026,0026 - Heico - Sigma - 7.avi,2,106008,106558,4240.32,4262.36,22.04,01:10:40.320,01:11:02.360,551,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0026,0026 - Heico - Sigma - 7.avi,3,135166,135550,5406.64,5422.04,15.4,01:30:06.640,01:30:22.040,385,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0026,0026 - Heico - Sigma - 7.avi,4,162447,162857,6497.88,6514.32,16.44,01:48:17.880,01:48:34.320,411,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0026,0026 - Heico - Sigma - 7.avi,5,179005,180062,7160.2,7202.52,42.32,01:59:20.200,02:00:02.520,1058,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0026,0026 - Heico - Sigma - 7.avi,6,184200,224505,7368.0,8980.24,1612.24,02:02:48.000,02:29:40.240,40306,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0026,0026 - Heico - Sigma - 7.avi,7,236957,238289,9478.28,9531.6,53.32,02:37:58.280,02:38:51.600,1333,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,1,1424,3592,56.96,143.72,86.76,00:00:56.960,00:02:23.720,2169,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,2,5305,5598,212.2,223.96,11.76,00:03:32.200,00:03:43.960,294,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,3,34580,34741,1383.2,1389.68,6.48,00:23:03.200,00:23:09.680,162,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,4,35894,36166,1435.76,1446.68,10.92,00:23:55.760,00:24:06.680,273,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,5,76803,119032,3072.12,4761.32,1689.2,00:51:12.120,01:19:21.320,42230,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,6,136180,136435,5447.2,5457.44,10.24,01:30:47.200,01:30:57.440,256,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,7,136515,136760,5460.6,5470.44,9.84,01:31:00.600,01:31:10.440,246,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,8,136865,137086,5474.6,5483.48,8.88,01:31:14.600,01:31:23.480,222,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,9,138468,138680,5538.72,5547.24,8.52,01:32:18.720,01:32:27.240,213,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,10,154029,154743,6161.16,6189.76,28.6,01:42:41.160,01:43:09.760,715,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,11,155735,155951,6229.4,6238.08,8.68,01:43:49.400,01:43:58.080,217,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,12,170430,171242,6817.2,6849.72,32.52,01:53:37.200,01:54:09.720,813,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,13,171760,173589,6870.4,6943.6,73.2,01:54:30.400,01:55:43.600,1830,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,14,175565,176292,7022.6,7051.72,29.12,01:57:02.600,01:57:31.720,728,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,15,188701,189011,7548.04,7560.48,12.44,02:05:48.040,02:06:00.480,311,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,16,193088,193299,7723.52,7732.0,8.48,02:08:43.520,02:08:52.000,212,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,17,193457,193676,7738.28,7747.08,8.8,02:08:58.280,02:09:07.080,220,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0027,0027 - Heico - Sigma - 8.avi,18,204905,205339,8196.2,8213.6,17.4,02:16:36.200,02:16:53.600,435,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,1,2902,4698,116.08,187.96,71.88,00:01:56.080,00:03:07.960,1797,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,2,4812,5000,192.48,200.04,7.56,00:03:12.480,00:03:20.040,189,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,3,5332,5945,213.28,237.84,24.56,00:03:33.280,00:03:57.840,614,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,4,6263,6645,250.52,265.84,15.32,00:04:10.520,00:04:25.840,383,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,5,81620,82270,3264.8,3290.84,26.04,00:54:24.800,00:54:50.840,651,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,6,83668,84298,3346.72,3371.96,25.24,00:55:46.720,00:56:11.960,631,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,7,91175,92008,3647.0,3680.36,33.36,01:00:47.000,01:01:20.360,834,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,8,98941,99826,3957.64,3993.08,35.44,01:05:57.640,01:06:33.080,886,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,9,99972,100157,3998.88,4006.32,7.44,01:06:38.880,01:06:46.320,186,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,10,106430,107644,4257.2,4305.8,48.6,01:10:57.200,01:11:45.800,1215,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,11,108850,109555,4354.0,4382.24,28.24,01:12:34.000,01:13:02.240,706,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,12,111863,113239,4474.52,4529.6,55.08,01:14:34.520,01:15:29.600,1377,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,13,113515,114657,4540.6,4586.32,45.72,01:15:40.600,01:16:26.320,1143,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,14,129364,146585,5174.56,5863.44,688.88,01:26:14.560,01:37:43.440,17222,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,15,157535,157730,6301.4,6309.24,7.84,01:45:01.400,01:45:09.240,196,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,16,164683,165551,6587.32,6622.08,34.76,01:49:47.320,01:50:22.080,869,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,17,178354,178553,7134.16,7142.16,8.0,01:58:54.160,01:59:02.160,200,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,18,183895,184094,7355.8,7363.8,8.0,02:02:35.800,02:02:43.800,200,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,19,188618,234529,7544.72,9381.2,1836.48,02:05:44.720,02:36:21.200,45912,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,20,284965,286101,11398.6,11444.08,45.48,03:09:58.600,03:10:44.080,1137,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,21,315904,316225,12636.16,12649.04,12.88,03:30:36.160,03:30:49.040,322,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,22,319567,364856,12782.68,14594.28,1811.6,03:33:02.680,04:03:14.280,45290,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0028,0028 - Heico - Sigma - 9.avi,23,381302,381504,15252.08,15260.2,8.12,04:14:12.080,04:14:20.200,203,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0029,0029 - Heico - Sigma - 10.avi,1,446,620,17.84,24.84,7.0,00:00:17.840,00:00:24.840,175,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0029,0029 - Heico - Sigma - 10.avi,2,1218,1586,48.72,63.48,14.76,00:00:48.720,00:01:03.480,369,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0029,0029 - Heico - Sigma - 10.avi,3,168562,217347,6742.48,8693.92,1951.44,01:52:22.480,02:24:53.920,48786,1.0,high,00c9758b75e5aaadf2d346b15f90db76 +0029,0029 - Heico - Sigma - 10.avi,4,222032,222368,8881.28,8894.76,13.48,02:28:01.280,02:28:14.760,337,1.0,high,00c9758b75e5aaadf2d346b15f90db76 diff --git a/pipelines/Tr006_V2/annotations/blue_screen/metadata/provenance.json b/pipelines/Tr006_V2/annotations/blue_screen/metadata/provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..df37fe89a55815da03c1ed0cfeb1770950427731 --- /dev/null +++ b/pipelines/Tr006_V2/annotations/blue_screen/metadata/provenance.json @@ -0,0 +1,43 @@ +{ + "source_dataset": "heico-focus-vqa", + "source_videos_included": false, + "run_id": "debug015", + "run_name": "blue_screen_coarse_to_fine_20260825", + "task_type": "debug", + "status": "completed", + "purpose": "Detect frozen blue-screen intervals in all 30 HeiCo endoscopy videos with 1-second coarse hashing and native-FPS exact boundary refinement.", + "code_commit": "a23500eae96b97f4b9aff3d00d08c3b8e7482aba", + "entrypoint": "data_preprocessing/process_blue_screen_dataset.py", + "parameters": { + "sample_seconds": 1.0, + "min_static_seconds": 6.0, + "blue_ratio": 0.45, + "pad_seconds": 5.0, + "coarse_size": 96, + "fine_size": 160, + "fine_resolution": "native FPS (25 FPS for inspected dataset)" + }, + "precision": "uint8 decoded frames", + "started_at": "2026-08-25T00:22:36+08:00", + "finished_at": "2026-08-25T02:54:37+08:00", + "results_summary": { + "video_count": 30, + "completed_video_count": 30, + "failed_video_count": 0, + "total_frames": 8681449, + "total_decoded_frames": 8681449, + "interval_count": 323, + "total_video_seconds": 347257.96, + "total_blue_screen_seconds": 82209.52, + "elapsed_seconds": 9119.6, + "validation": "PASS" + }, + "issues": [], + "acceptance": { + "all_30_videos_accounted_for": true, + "no_source_video_modification": true, + "frame_level_boundaries": true, + "complete_log": true, + "aggregate_counts_match_per_video_results": true + } +} diff --git a/pipelines/Tr006_V2/annotations/blue_screen/metadata/summary.json b/pipelines/Tr006_V2/annotations/blue_screen/metadata/summary.json new file mode 100644 index 0000000000000000000000000000000000000000..378d988cc9be7c777be471f4103cda688303936a --- /dev/null +++ b/pipelines/Tr006_V2/annotations/blue_screen/metadata/summary.json @@ -0,0 +1,23 @@ +{ + "status": "completed", + "started_at": "2026-08-25T00:22:37.492532+08:00", + "finished_at": "2026-08-25T02:54:37.111581+08:00", + "elapsed_seconds": 9119.599936511368, + "host": "farm02", + "config": { + "sample_seconds": 1.0, + "min_static_seconds": 6.0, + "blue_ratio": 0.45, + "pad_seconds": 5.0, + "coarse_size": 96, + "fine_size": 160, + "definition": "exact repeated decoded ROI frames with blue ratio above threshold" + }, + "video_count": 30, + "completed_video_count": 30, + "failed_video_count": 0, + "videos_with_blue_screen": 30, + "interval_count": 323, + "total_video_seconds": 347257.9600000001, + "total_blue_screen_seconds": 82209.52000000003 +} diff --git a/pipelines/Tr006_V2/annotations/blue_screen/metadata/validation.json b/pipelines/Tr006_V2/annotations/blue_screen/metadata/validation.json new file mode 100644 index 0000000000000000000000000000000000000000..45cf01f40fc41546e5bdfd7dc99d3416d715f3a9 --- /dev/null +++ b/pipelines/Tr006_V2/annotations/blue_screen/metadata/validation.json @@ -0,0 +1,30 @@ +{ + "status": "PASS", + "validated_at": "2026-08-25T02:54:37+08:00", + "source_video_count": 30, + "per_video_json_count": 30, + "videos_jsonl_count": 30, + "completed_video_count": 30, + "failed_video_count": 0, + "interval_count": 323, + "csv_interval_count": 323, + "evidence_jpeg_count": 1291, + "total_frames": 8681449, + "total_decoded_frames": 8681449, + "total_video_seconds": 347257.96, + "total_blue_screen_seconds": 82209.52, + "checks": { + "source_paths_exactly_covered": true, + "all_videos_completed": true, + "all_frames_decoded": true, + "intervals_ordered_non_overlapping_and_in_bounds": true, + "interval_duration_matches_frame_bounds": true, + "all_intervals_meet_hash_and_blue_criteria": true, + "aggregate_csv_matches_per_video_results": true, + "all_non_boundary_evidence_frames_present": true, + "run_log_exit_code_zero": true + }, + "notes": [ + "Video 0022 has one interval ending at the final source frame, so an after-boundary evidence frame does not exist by definition." + ] +} diff --git a/pipelines/Tr006_V2/annotations/blue_screen/metadata/videos.jsonl b/pipelines/Tr006_V2/annotations/blue_screen/metadata/videos.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..aceed5fbe83c73e695a8b1c251a8e093d620cc99 --- /dev/null +++ b/pipelines/Tr006_V2/annotations/blue_screen/metadata/videos.jsonl @@ -0,0 +1,30 @@ +{"video_id": "0000", "status": "completed", "started_at": "2026-08-25T00:22:37.494313+08:00", "intervals": [{"start_frame": 5524, "end_frame": 5938, "start_seconds": 220.96, "end_seconds": 237.56, "duration_seconds": 16.6, "start_time": "00:03:40.960", "end_time": "00:03:57.560", "frame_count": 415, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0000/interval_001_before_f5523.jpg", "start": "evidence/0000/interval_001_start_f5524.jpg", "end": "evidence/0000/interval_001_end_f5938.jpg", "after": "evidence/0000/interval_001_after_f5939.jpg"}}, {"start_frame": 9296, "end_frame": 9548, "start_seconds": 371.84, "end_seconds": 381.96, "duration_seconds": 10.12, "start_time": "00:06:11.840", "end_time": "00:06:21.960", "frame_count": 253, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0000/interval_002_before_f9295.jpg", "start": "evidence/0000/interval_002_start_f9296.jpg", "end": "evidence/0000/interval_002_end_f9548.jpg", "after": "evidence/0000/interval_002_after_f9549.jpg"}}, {"start_frame": 10249, "end_frame": 10494, "start_seconds": 409.96, "end_seconds": 419.8, "duration_seconds": 9.84, "start_time": "00:06:49.960", "end_time": "00:06:59.800", "frame_count": 246, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0000/interval_003_before_f10248.jpg", "start": "evidence/0000/interval_003_start_f10249.jpg", "end": "evidence/0000/interval_003_end_f10494.jpg", "after": "evidence/0000/interval_003_after_f10495.jpg"}}, {"start_frame": 105966, "end_frame": 106309, "start_seconds": 4238.64, "end_seconds": 4252.4, "duration_seconds": 13.76, "start_time": "01:10:38.640", "end_time": "01:10:52.400", "frame_count": 344, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0000/interval_004_before_f105965.jpg", "start": "evidence/0000/interval_004_start_f105966.jpg", "end": "evidence/0000/interval_004_end_f106309.jpg", "after": "evidence/0000/interval_004_after_f106310.jpg"}}, {"start_frame": 153266, "end_frame": 153973, "start_seconds": 6130.64, "end_seconds": 6158.96, "duration_seconds": 28.32, "start_time": "01:42:10.640", "end_time": "01:42:38.960", "frame_count": 708, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0000/interval_005_before_f153265.jpg", "start": "evidence/0000/interval_005_start_f153266.jpg", "end": "evidence/0000/interval_005_end_f153973.jpg", "after": "evidence/0000/interval_005_after_f153974.jpg"}}, {"start_frame": 216719, "end_frame": 217948, "start_seconds": 8668.76, "end_seconds": 8717.96, "duration_seconds": 49.2, "start_time": "02:24:28.760", "end_time": "02:25:17.960", "frame_count": 1230, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0000/interval_006_before_f216718.jpg", "start": "evidence/0000/interval_006_start_f216719.jpg", "end": "evidence/0000/interval_006_end_f217948.jpg", "after": "evidence/0000/interval_006_after_f217949.jpg"}}, {"start_frame": 232574, "end_frame": 311655, "start_seconds": 9302.96, "end_seconds": 12466.24, "duration_seconds": 3163.28, "start_time": "02:35:02.960", "end_time": "03:27:46.240", "frame_count": 79082, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0000/interval_007_before_f232573.jpg", "start": "evidence/0000/interval_007_start_f232574.jpg", "end": "evidence/0000/interval_007_end_f311655.jpg", "after": "evidence/0000/interval_007_after_f311656.jpg"}}], "fps": 25.0, "frame_count": 371603, "width": 960, "height": 540, "duration_seconds": 14864.12, "decoded_frames": 371603, "coarse_samples": 14865, "coarse_candidate_count": 7, "interval_count": 7, "blue_screen": true, "blue_screen_seconds": 3291.1200000000003, "blue_screen_ratio": 0.22141371302169252, "finished_at": "2026-08-25T00:28:02.142773+08:00", "elapsed_seconds": 324.64850525557995, "video_filename": "0000 - Heico - Prokto - 1.avi"} +{"video_id": "0001", "status": "completed", "started_at": "2026-08-25T00:28:02.216867+08:00", "intervals": [{"start_frame": 110844, "end_frame": 111882, "start_seconds": 4433.76, "end_seconds": 4475.32, "duration_seconds": 41.56, "start_time": "01:13:53.760", "end_time": "01:14:35.320", "frame_count": 1039, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0001/interval_001_before_f110843.jpg", "start": "evidence/0001/interval_001_start_f110844.jpg", "end": "evidence/0001/interval_001_end_f111882.jpg", "after": "evidence/0001/interval_001_after_f111883.jpg"}}, {"start_frame": 112792, "end_frame": 113563, "start_seconds": 4511.68, "end_seconds": 4542.56, "duration_seconds": 30.88, "start_time": "01:15:11.680", "end_time": "01:15:42.560", "frame_count": 772, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0001/interval_002_before_f112791.jpg", "start": "evidence/0001/interval_002_start_f112792.jpg", "end": "evidence/0001/interval_002_end_f113563.jpg", "after": "evidence/0001/interval_002_after_f113564.jpg"}}, {"start_frame": 114462, "end_frame": 115183, "start_seconds": 4578.48, "end_seconds": 4607.36, "duration_seconds": 28.88, "start_time": "01:16:18.480", "end_time": "01:16:47.360", "frame_count": 722, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0001/interval_003_before_f114461.jpg", "start": "evidence/0001/interval_003_start_f114462.jpg", "end": "evidence/0001/interval_003_end_f115183.jpg", "after": "evidence/0001/interval_003_after_f115184.jpg"}}, {"start_frame": 137203, "end_frame": 138843, "start_seconds": 5488.12, "end_seconds": 5553.76, "duration_seconds": 65.64, "start_time": "01:31:28.120", "end_time": "01:32:33.760", "frame_count": 1641, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0001/interval_004_before_f137202.jpg", "start": "evidence/0001/interval_004_start_f137203.jpg", "end": "evidence/0001/interval_004_end_f138843.jpg", "after": "evidence/0001/interval_004_after_f138844.jpg"}}, {"start_frame": 147830, "end_frame": 148183, "start_seconds": 5913.2, "end_seconds": 5927.36, "duration_seconds": 14.16, "start_time": "01:38:33.200", "end_time": "01:38:47.360", "frame_count": 354, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0001/interval_005_before_f147829.jpg", "start": "evidence/0001/interval_005_start_f147830.jpg", "end": "evidence/0001/interval_005_end_f148183.jpg", "after": "evidence/0001/interval_005_after_f148184.jpg"}}, {"start_frame": 154539, "end_frame": 155178, "start_seconds": 6181.56, "end_seconds": 6207.16, "duration_seconds": 25.6, "start_time": "01:43:01.560", "end_time": "01:43:27.160", "frame_count": 640, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0001/interval_006_before_f154538.jpg", "start": "evidence/0001/interval_006_start_f154539.jpg", "end": "evidence/0001/interval_006_end_f155178.jpg", "after": "evidence/0001/interval_006_after_f155179.jpg"}}, {"start_frame": 222864, "end_frame": 223263, "start_seconds": 8914.56, "end_seconds": 8930.56, "duration_seconds": 16.0, "start_time": "02:28:34.560", "end_time": "02:28:50.560", "frame_count": 400, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0001/interval_007_before_f222863.jpg", "start": "evidence/0001/interval_007_start_f222864.jpg", "end": "evidence/0001/interval_007_end_f223263.jpg", "after": "evidence/0001/interval_007_after_f223264.jpg"}}, {"start_frame": 237013, "end_frame": 316761, "start_seconds": 9480.52, "end_seconds": 12670.48, "duration_seconds": 3189.96, "start_time": "02:38:00.520", "end_time": "03:31:10.480", "frame_count": 79749, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0001/interval_008_before_f237012.jpg", "start": "evidence/0001/interval_008_start_f237013.jpg", "end": "evidence/0001/interval_008_end_f316761.jpg", "after": "evidence/0001/interval_008_after_f316762.jpg"}}], "fps": 25.0, "frame_count": 350519, "width": 960, "height": 540, "duration_seconds": 14020.76, "decoded_frames": 350519, "coarse_samples": 14021, "coarse_candidate_count": 8, "interval_count": 8, "blue_screen": true, "blue_screen_seconds": 3412.68, "blue_screen_ratio": 0.243401926857032, "finished_at": "2026-08-25T00:35:00.028984+08:00", "elapsed_seconds": 417.8121481779963, "video_filename": "0001 - Heico - Prokto - 2.avi"} +{"video_id": "0002", "status": "completed", "started_at": "2026-08-25T00:35:00.043642+08:00", "intervals": [{"start_frame": 5009, "end_frame": 5212, "start_seconds": 200.36, "end_seconds": 208.52, "duration_seconds": 8.16, "start_time": "00:03:20.360", "end_time": "00:03:28.520", "frame_count": 204, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_001_before_f5008.jpg", "start": "evidence/0002/interval_001_start_f5009.jpg", "end": "evidence/0002/interval_001_end_f5212.jpg", "after": "evidence/0002/interval_001_after_f5213.jpg"}}, {"start_frame": 5326, "end_frame": 5704, "start_seconds": 213.04, "end_seconds": 228.2, "duration_seconds": 15.16, "start_time": "00:03:33.040", "end_time": "00:03:48.200", "frame_count": 379, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_002_before_f5325.jpg", "start": "evidence/0002/interval_002_start_f5326.jpg", "end": "evidence/0002/interval_002_end_f5704.jpg", "after": "evidence/0002/interval_002_after_f5705.jpg"}}, {"start_frame": 26052, "end_frame": 27664, "start_seconds": 1042.08, "end_seconds": 1106.6, "duration_seconds": 64.52, "start_time": "00:17:22.080", "end_time": "00:18:26.600", "frame_count": 1613, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_003_before_f26051.jpg", "start": "evidence/0002/interval_003_start_f26052.jpg", "end": "evidence/0002/interval_003_end_f27664.jpg", "after": "evidence/0002/interval_003_after_f27665.jpg"}}, {"start_frame": 115654, "end_frame": 115924, "start_seconds": 4626.16, "end_seconds": 4637.0, "duration_seconds": 10.84, "start_time": "01:17:06.160", "end_time": "01:17:17.000", "frame_count": 271, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_004_before_f115653.jpg", "start": "evidence/0002/interval_004_start_f115654.jpg", "end": "evidence/0002/interval_004_end_f115924.jpg", "after": "evidence/0002/interval_004_after_f115925.jpg"}}, {"start_frame": 136596, "end_frame": 136848, "start_seconds": 5463.84, "end_seconds": 5473.96, "duration_seconds": 10.12, "start_time": "01:31:03.840", "end_time": "01:31:13.960", "frame_count": 253, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_005_before_f136595.jpg", "start": "evidence/0002/interval_005_start_f136596.jpg", "end": "evidence/0002/interval_005_end_f136848.jpg", "after": "evidence/0002/interval_005_after_f136849.jpg"}}, {"start_frame": 144722, "end_frame": 144905, "start_seconds": 5788.88, "end_seconds": 5796.24, "duration_seconds": 7.36, "start_time": "01:36:28.880", "end_time": "01:36:36.240", "frame_count": 184, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_006_before_f144721.jpg", "start": "evidence/0002/interval_006_start_f144722.jpg", "end": "evidence/0002/interval_006_end_f144905.jpg", "after": "evidence/0002/interval_006_after_f144906.jpg"}}, {"start_frame": 159926, "end_frame": 217858, "start_seconds": 6397.04, "end_seconds": 8714.36, "duration_seconds": 2317.32, "start_time": "01:46:37.040", "end_time": "02:25:14.360", "frame_count": 57933, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_007_before_f159925.jpg", "start": "evidence/0002/interval_007_start_f159926.jpg", "end": "evidence/0002/interval_007_end_f217858.jpg", "after": "evidence/0002/interval_007_after_f217859.jpg"}}, {"start_frame": 227798, "end_frame": 229000, "start_seconds": 9111.92, "end_seconds": 9160.04, "duration_seconds": 48.12, "start_time": "02:31:51.920", "end_time": "02:32:40.040", "frame_count": 1203, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_008_before_f227797.jpg", "start": "evidence/0002/interval_008_start_f227798.jpg", "end": "evidence/0002/interval_008_end_f229000.jpg", "after": "evidence/0002/interval_008_after_f229001.jpg"}}, {"start_frame": 248690, "end_frame": 251825, "start_seconds": 9947.6, "end_seconds": 10073.04, "duration_seconds": 125.44, "start_time": "02:45:47.600", "end_time": "02:47:53.040", "frame_count": 3136, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_009_before_f248689.jpg", "start": "evidence/0002/interval_009_start_f248690.jpg", "end": "evidence/0002/interval_009_end_f251825.jpg", "after": "evidence/0002/interval_009_after_f251826.jpg"}}, {"start_frame": 252911, "end_frame": 254146, "start_seconds": 10116.44, "end_seconds": 10165.88, "duration_seconds": 49.44, "start_time": "02:48:36.440", "end_time": "02:49:25.880", "frame_count": 1236, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_010_before_f252910.jpg", "start": "evidence/0002/interval_010_start_f252911.jpg", "end": "evidence/0002/interval_010_end_f254146.jpg", "after": "evidence/0002/interval_010_after_f254147.jpg"}}, {"start_frame": 255848, "end_frame": 256281, "start_seconds": 10233.92, "end_seconds": 10251.28, "duration_seconds": 17.36, "start_time": "02:50:33.920", "end_time": "02:50:51.280", "frame_count": 434, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0002/interval_011_before_f255847.jpg", "start": "evidence/0002/interval_011_start_f255848.jpg", "end": "evidence/0002/interval_011_end_f256281.jpg", "after": "evidence/0002/interval_011_after_f256282.jpg"}}], "fps": 25.0, "frame_count": 278952, "width": 960, "height": 540, "duration_seconds": 11158.08, "decoded_frames": 278952, "coarse_samples": 11159, "coarse_candidate_count": 11, "interval_count": 11, "blue_screen": true, "blue_screen_seconds": 2673.84, "blue_screen_ratio": 0.2396326249677364, "finished_at": "2026-08-25T00:40:13.235731+08:00", "elapsed_seconds": 313.19211813248694, "video_filename": "0002 - Heico - Prokto - 3.avi"} +{"video_id": "0003", "status": "completed", "started_at": "2026-08-25T00:40:13.267265+08:00", "intervals": [{"start_frame": 643, "end_frame": 869, "start_seconds": 25.72, "end_seconds": 34.8, "duration_seconds": 9.08, "start_time": "00:00:25.720", "end_time": "00:00:34.800", "frame_count": 227, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_001_before_f642.jpg", "start": "evidence/0003/interval_001_start_f643.jpg", "end": "evidence/0003/interval_001_end_f869.jpg", "after": "evidence/0003/interval_001_after_f870.jpg"}}, {"start_frame": 1267, "end_frame": 1562, "start_seconds": 50.68, "end_seconds": 62.52, "duration_seconds": 11.84, "start_time": "00:00:50.680", "end_time": "00:01:02.520", "frame_count": 296, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_002_before_f1266.jpg", "start": "evidence/0003/interval_002_start_f1267.jpg", "end": "evidence/0003/interval_002_end_f1562.jpg", "after": "evidence/0003/interval_002_after_f1563.jpg"}}, {"start_frame": 11885, "end_frame": 12122, "start_seconds": 475.4, "end_seconds": 484.92, "duration_seconds": 9.52, "start_time": "00:07:55.400", "end_time": "00:08:04.920", "frame_count": 238, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_003_before_f11884.jpg", "start": "evidence/0003/interval_003_start_f11885.jpg", "end": "evidence/0003/interval_003_end_f12122.jpg", "after": "evidence/0003/interval_003_after_f12123.jpg"}}, {"start_frame": 82913, "end_frame": 83622, "start_seconds": 3316.52, "end_seconds": 3344.92, "duration_seconds": 28.4, "start_time": "00:55:16.520", "end_time": "00:55:44.920", "frame_count": 710, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_004_before_f82912.jpg", "start": "evidence/0003/interval_004_start_f82913.jpg", "end": "evidence/0003/interval_004_end_f83622.jpg", "after": "evidence/0003/interval_004_after_f83623.jpg"}}, {"start_frame": 127912, "end_frame": 128071, "start_seconds": 5116.48, "end_seconds": 5122.88, "duration_seconds": 6.4, "start_time": "01:25:16.480", "end_time": "01:25:22.880", "frame_count": 160, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_005_before_f127911.jpg", "start": "evidence/0003/interval_005_start_f127912.jpg", "end": "evidence/0003/interval_005_end_f128071.jpg", "after": "evidence/0003/interval_005_after_f128072.jpg"}}, {"start_frame": 141561, "end_frame": 200111, "start_seconds": 5662.44, "end_seconds": 8004.48, "duration_seconds": 2342.04, "start_time": "01:34:22.440", "end_time": "02:13:24.480", "frame_count": 58551, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_006_before_f141560.jpg", "start": "evidence/0003/interval_006_start_f141561.jpg", "end": "evidence/0003/interval_006_end_f200111.jpg", "after": "evidence/0003/interval_006_after_f200112.jpg"}}, {"start_frame": 205346, "end_frame": 206231, "start_seconds": 8213.84, "end_seconds": 8249.28, "duration_seconds": 35.44, "start_time": "02:16:53.840", "end_time": "02:17:29.280", "frame_count": 886, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_007_before_f205345.jpg", "start": "evidence/0003/interval_007_start_f205346.jpg", "end": "evidence/0003/interval_007_end_f206231.jpg", "after": "evidence/0003/interval_007_after_f206232.jpg"}}, {"start_frame": 233424, "end_frame": 237433, "start_seconds": 9336.96, "end_seconds": 9497.36, "duration_seconds": 160.4, "start_time": "02:35:36.960", "end_time": "02:38:17.360", "frame_count": 4010, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_008_before_f233423.jpg", "start": "evidence/0003/interval_008_start_f233424.jpg", "end": "evidence/0003/interval_008_end_f237433.jpg", "after": "evidence/0003/interval_008_after_f237434.jpg"}}, {"start_frame": 254178, "end_frame": 255182, "start_seconds": 10167.12, "end_seconds": 10207.32, "duration_seconds": 40.2, "start_time": "02:49:27.120", "end_time": "02:50:07.320", "frame_count": 1005, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_009_before_f254177.jpg", "start": "evidence/0003/interval_009_start_f254178.jpg", "end": "evidence/0003/interval_009_end_f255182.jpg", "after": "evidence/0003/interval_009_after_f255183.jpg"}}, {"start_frame": 255278, "end_frame": 255922, "start_seconds": 10211.12, "end_seconds": 10236.92, "duration_seconds": 25.8, "start_time": "02:50:11.120", "end_time": "02:50:36.920", "frame_count": 645, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0003/interval_010_before_f255277.jpg", "start": "evidence/0003/interval_010_start_f255278.jpg", "end": "evidence/0003/interval_010_end_f255922.jpg", "after": "evidence/0003/interval_010_after_f255923.jpg"}}], "fps": 25.0, "frame_count": 261521, "width": 960, "height": 540, "duration_seconds": 10460.84, "decoded_frames": 261521, "coarse_samples": 10461, "coarse_candidate_count": 10, "interval_count": 10, "blue_screen": true, "blue_screen_seconds": 2669.12, "blue_screen_ratio": 0.25515350583700735, "finished_at": "2026-08-25T00:44:46.406541+08:00", "elapsed_seconds": 273.1393105406314, "video_filename": "0003 - Heico - Prokto - 4.avi"} +{"video_id": "0004", "status": "completed", "started_at": "2026-08-25T00:44:46.419283+08:00", "intervals": [{"start_frame": 811, "end_frame": 1005, "start_seconds": 32.44, "end_seconds": 40.24, "duration_seconds": 7.8, "start_time": "00:00:32.440", "end_time": "00:00:40.240", "frame_count": 195, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_001_before_f810.jpg", "start": "evidence/0004/interval_001_start_f811.jpg", "end": "evidence/0004/interval_001_end_f1005.jpg", "after": "evidence/0004/interval_001_after_f1006.jpg"}}, {"start_frame": 2507, "end_frame": 2827, "start_seconds": 100.28, "end_seconds": 113.12, "duration_seconds": 12.84, "start_time": "00:01:40.280", "end_time": "00:01:53.120", "frame_count": 321, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_002_before_f2506.jpg", "start": "evidence/0004/interval_002_start_f2507.jpg", "end": "evidence/0004/interval_002_end_f2827.jpg", "after": "evidence/0004/interval_002_after_f2828.jpg"}}, {"start_frame": 2911, "end_frame": 3179, "start_seconds": 116.44, "end_seconds": 127.2, "duration_seconds": 10.76, "start_time": "00:01:56.440", "end_time": "00:02:07.200", "frame_count": 269, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_003_before_f2910.jpg", "start": "evidence/0004/interval_003_start_f2911.jpg", "end": "evidence/0004/interval_003_end_f3179.jpg", "after": "evidence/0004/interval_003_after_f3180.jpg"}}, {"start_frame": 5065, "end_frame": 5270, "start_seconds": 202.6, "end_seconds": 210.84, "duration_seconds": 8.24, "start_time": "00:03:22.600", "end_time": "00:03:30.840", "frame_count": 206, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_004_before_f5064.jpg", "start": "evidence/0004/interval_004_start_f5065.jpg", "end": "evidence/0004/interval_004_end_f5270.jpg", "after": "evidence/0004/interval_004_after_f5271.jpg"}}, {"start_frame": 5487, "end_frame": 5670, "start_seconds": 219.48, "end_seconds": 226.84, "duration_seconds": 7.36, "start_time": "00:03:39.480", "end_time": "00:03:46.840", "frame_count": 184, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_005_before_f5486.jpg", "start": "evidence/0004/interval_005_start_f5487.jpg", "end": "evidence/0004/interval_005_end_f5670.jpg", "after": "evidence/0004/interval_005_after_f5671.jpg"}}, {"start_frame": 153423, "end_frame": 153582, "start_seconds": 6136.92, "end_seconds": 6143.32, "duration_seconds": 6.4, "start_time": "01:42:16.920", "end_time": "01:42:23.320", "frame_count": 160, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_006_before_f153422.jpg", "start": "evidence/0004/interval_006_start_f153423.jpg", "end": "evidence/0004/interval_006_end_f153582.jpg", "after": "evidence/0004/interval_006_after_f153583.jpg"}}, {"start_frame": 166711, "end_frame": 237718, "start_seconds": 6668.44, "end_seconds": 9508.76, "duration_seconds": 2840.32, "start_time": "01:51:08.440", "end_time": "02:38:28.760", "frame_count": 71008, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_007_before_f166710.jpg", "start": "evidence/0004/interval_007_start_f166711.jpg", "end": "evidence/0004/interval_007_end_f237718.jpg", "after": "evidence/0004/interval_007_after_f237719.jpg"}}, {"start_frame": 274728, "end_frame": 274887, "start_seconds": 10989.12, "end_seconds": 10995.52, "duration_seconds": 6.4, "start_time": "03:03:09.120", "end_time": "03:03:15.520", "frame_count": 160, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_008_before_f274727.jpg", "start": "evidence/0004/interval_008_start_f274728.jpg", "end": "evidence/0004/interval_008_end_f274887.jpg", "after": "evidence/0004/interval_008_after_f274888.jpg"}}, {"start_frame": 279748, "end_frame": 280329, "start_seconds": 11189.92, "end_seconds": 11213.2, "duration_seconds": 23.28, "start_time": "03:06:29.920", "end_time": "03:06:53.200", "frame_count": 582, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_009_before_f279747.jpg", "start": "evidence/0004/interval_009_start_f279748.jpg", "end": "evidence/0004/interval_009_end_f280329.jpg", "after": "evidence/0004/interval_009_after_f280330.jpg"}}, {"start_frame": 281024, "end_frame": 281672, "start_seconds": 11240.96, "end_seconds": 11266.92, "duration_seconds": 25.96, "start_time": "03:07:20.960", "end_time": "03:07:46.920", "frame_count": 649, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_010_before_f281023.jpg", "start": "evidence/0004/interval_010_start_f281024.jpg", "end": "evidence/0004/interval_010_end_f281672.jpg", "after": "evidence/0004/interval_010_after_f281673.jpg"}}, {"start_frame": 295193, "end_frame": 298502, "start_seconds": 11807.72, "end_seconds": 11940.12, "duration_seconds": 132.4, "start_time": "03:16:47.720", "end_time": "03:19:00.120", "frame_count": 3310, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_011_before_f295192.jpg", "start": "evidence/0004/interval_011_start_f295193.jpg", "end": "evidence/0004/interval_011_end_f298502.jpg", "after": "evidence/0004/interval_011_after_f298503.jpg"}}, {"start_frame": 299039, "end_frame": 299984, "start_seconds": 11961.56, "end_seconds": 11999.4, "duration_seconds": 37.84, "start_time": "03:19:21.560", "end_time": "03:19:59.400", "frame_count": 946, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_012_before_f299038.jpg", "start": "evidence/0004/interval_012_start_f299039.jpg", "end": "evidence/0004/interval_012_end_f299984.jpg", "after": "evidence/0004/interval_012_after_f299985.jpg"}}, {"start_frame": 301545, "end_frame": 302528, "start_seconds": 12061.8, "end_seconds": 12101.16, "duration_seconds": 39.36, "start_time": "03:21:01.800", "end_time": "03:21:41.160", "frame_count": 984, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_013_before_f301544.jpg", "start": "evidence/0004/interval_013_start_f301545.jpg", "end": "evidence/0004/interval_013_end_f302528.jpg", "after": "evidence/0004/interval_013_after_f302529.jpg"}}, {"start_frame": 304448, "end_frame": 305829, "start_seconds": 12177.92, "end_seconds": 12233.2, "duration_seconds": 55.28, "start_time": "03:22:57.920", "end_time": "03:23:53.200", "frame_count": 1382, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_014_before_f304447.jpg", "start": "evidence/0004/interval_014_start_f304448.jpg", "end": "evidence/0004/interval_014_end_f305829.jpg", "after": "evidence/0004/interval_014_after_f305830.jpg"}}, {"start_frame": 306273, "end_frame": 312447, "start_seconds": 12250.92, "end_seconds": 12497.92, "duration_seconds": 247.0, "start_time": "03:24:10.920", "end_time": "03:28:17.920", "frame_count": 6175, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_015_before_f306272.jpg", "start": "evidence/0004/interval_015_start_f306273.jpg", "end": "evidence/0004/interval_015_end_f312447.jpg", "after": "evidence/0004/interval_015_after_f312448.jpg"}}, {"start_frame": 313573, "end_frame": 313830, "start_seconds": 12542.92, "end_seconds": 12553.24, "duration_seconds": 10.32, "start_time": "03:29:02.920", "end_time": "03:29:13.240", "frame_count": 258, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_016_before_f313572.jpg", "start": "evidence/0004/interval_016_start_f313573.jpg", "end": "evidence/0004/interval_016_end_f313830.jpg", "after": "evidence/0004/interval_016_after_f313831.jpg"}}, {"start_frame": 315328, "end_frame": 315503, "start_seconds": 12613.12, "end_seconds": 12620.16, "duration_seconds": 7.04, "start_time": "03:30:13.120", "end_time": "03:30:20.160", "frame_count": 176, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0004/interval_017_before_f315327.jpg", "start": "evidence/0004/interval_017_start_f315328.jpg", "end": "evidence/0004/interval_017_end_f315503.jpg", "after": "evidence/0004/interval_017_after_f315504.jpg"}}], "fps": 25.0, "frame_count": 317924, "width": 960, "height": 540, "duration_seconds": 12716.96, "decoded_frames": 317924, "coarse_samples": 12717, "coarse_candidate_count": 18, "interval_count": 17, "blue_screen": true, "blue_screen_seconds": 3478.6000000000013, "blue_screen_ratio": 0.27354021715881793, "finished_at": "2026-08-25T00:50:56.335007+08:00", "elapsed_seconds": 369.9157545976341, "video_filename": "0004 - Heico - Prokto - 5.avi"} +{"video_id": "0005", "status": "completed", "started_at": "2026-08-25T00:50:56.368063+08:00", "intervals": [{"start_frame": 63, "end_frame": 977, "start_seconds": 2.52, "end_seconds": 39.12, "duration_seconds": 36.6, "start_time": "00:00:02.520", "end_time": "00:00:39.120", "frame_count": 915, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_001_before_f62.jpg", "start": "evidence/0005/interval_001_start_f63.jpg", "end": "evidence/0005/interval_001_end_f977.jpg", "after": "evidence/0005/interval_001_after_f978.jpg"}}, {"start_frame": 2600, "end_frame": 2834, "start_seconds": 104.0, "end_seconds": 113.4, "duration_seconds": 9.4, "start_time": "00:01:44.000", "end_time": "00:01:53.400", "frame_count": 235, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_002_before_f2599.jpg", "start": "evidence/0005/interval_002_start_f2600.jpg", "end": "evidence/0005/interval_002_end_f2834.jpg", "after": "evidence/0005/interval_002_after_f2835.jpg"}}, {"start_frame": 88472, "end_frame": 88932, "start_seconds": 3538.88, "end_seconds": 3557.32, "duration_seconds": 18.44, "start_time": "00:58:58.880", "end_time": "00:59:17.320", "frame_count": 461, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_003_before_f88471.jpg", "start": "evidence/0005/interval_003_start_f88472.jpg", "end": "evidence/0005/interval_003_end_f88932.jpg", "after": "evidence/0005/interval_003_after_f88933.jpg"}}, {"start_frame": 89225, "end_frame": 89747, "start_seconds": 3569.0, "end_seconds": 3589.92, "duration_seconds": 20.92, "start_time": "00:59:29.000", "end_time": "00:59:49.920", "frame_count": 523, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_004_before_f89224.jpg", "start": "evidence/0005/interval_004_start_f89225.jpg", "end": "evidence/0005/interval_004_end_f89747.jpg", "after": "evidence/0005/interval_004_after_f89748.jpg"}}, {"start_frame": 89971, "end_frame": 90170, "start_seconds": 3598.84, "end_seconds": 3606.84, "duration_seconds": 8.0, "start_time": "00:59:58.840", "end_time": "01:00:06.840", "frame_count": 200, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_005_before_f89970.jpg", "start": "evidence/0005/interval_005_start_f89971.jpg", "end": "evidence/0005/interval_005_end_f90170.jpg", "after": "evidence/0005/interval_005_after_f90171.jpg"}}, {"start_frame": 117330, "end_frame": 166097, "start_seconds": 4693.2, "end_seconds": 6643.92, "duration_seconds": 1950.72, "start_time": "01:18:13.200", "end_time": "01:50:43.920", "frame_count": 48768, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_006_before_f117329.jpg", "start": "evidence/0005/interval_006_start_f117330.jpg", "end": "evidence/0005/interval_006_end_f166097.jpg", "after": "evidence/0005/interval_006_after_f166098.jpg"}}, {"start_frame": 168608, "end_frame": 168844, "start_seconds": 6744.32, "end_seconds": 6753.8, "duration_seconds": 9.48, "start_time": "01:52:24.320", "end_time": "01:52:33.800", "frame_count": 237, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_007_before_f168607.jpg", "start": "evidence/0005/interval_007_start_f168608.jpg", "end": "evidence/0005/interval_007_end_f168844.jpg", "after": "evidence/0005/interval_007_after_f168845.jpg"}}, {"start_frame": 184972, "end_frame": 185273, "start_seconds": 7398.88, "end_seconds": 7410.96, "duration_seconds": 12.08, "start_time": "02:03:18.880", "end_time": "02:03:30.960", "frame_count": 302, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_008_before_f184971.jpg", "start": "evidence/0005/interval_008_start_f184972.jpg", "end": "evidence/0005/interval_008_end_f185273.jpg", "after": "evidence/0005/interval_008_after_f185274.jpg"}}, {"start_frame": 199269, "end_frame": 199539, "start_seconds": 7970.76, "end_seconds": 7981.6, "duration_seconds": 10.84, "start_time": "02:12:50.760", "end_time": "02:13:01.600", "frame_count": 271, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0005/interval_009_before_f199268.jpg", "start": "evidence/0005/interval_009_start_f199269.jpg", "end": "evidence/0005/interval_009_end_f199539.jpg", "after": "evidence/0005/interval_009_after_f199540.jpg"}}], "fps": 25.0, "frame_count": 204090, "width": 960, "height": 540, "duration_seconds": 8163.6, "decoded_frames": 204090, "coarse_samples": 8164, "coarse_candidate_count": 9, "interval_count": 9, "blue_screen": true, "blue_screen_seconds": 2076.48, "blue_screen_ratio": 0.2543583713067764, "finished_at": "2026-08-25T00:54:44.272649+08:00", "elapsed_seconds": 227.90461849980056, "video_filename": "0005 - Heico - Prokto - 6.avi"} +{"video_id": "0006", "status": "completed", "started_at": "2026-08-25T00:54:44.283963+08:00", "intervals": [{"start_frame": 54370, "end_frame": 54766, "start_seconds": 2174.8, "end_seconds": 2190.68, "duration_seconds": 15.88, "start_time": "00:36:14.800", "end_time": "00:36:30.680", "frame_count": 397, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_001_before_f54369.jpg", "start": "evidence/0006/interval_001_start_f54370.jpg", "end": "evidence/0006/interval_001_end_f54766.jpg", "after": "evidence/0006/interval_001_after_f54767.jpg"}}, {"start_frame": 68517, "end_frame": 68799, "start_seconds": 2740.68, "end_seconds": 2752.0, "duration_seconds": 11.32, "start_time": "00:45:40.680", "end_time": "00:45:52.000", "frame_count": 283, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_002_before_f68516.jpg", "start": "evidence/0006/interval_002_start_f68517.jpg", "end": "evidence/0006/interval_002_end_f68799.jpg", "after": "evidence/0006/interval_002_after_f68800.jpg"}}, {"start_frame": 131005, "end_frame": 192543, "start_seconds": 5240.2, "end_seconds": 7701.76, "duration_seconds": 2461.56, "start_time": "01:27:20.200", "end_time": "02:08:21.760", "frame_count": 61539, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_003_before_f131004.jpg", "start": "evidence/0006/interval_003_start_f131005.jpg", "end": "evidence/0006/interval_003_end_f192543.jpg", "after": "evidence/0006/interval_003_after_f192544.jpg"}}, {"start_frame": 192597, "end_frame": 192916, "start_seconds": 7703.88, "end_seconds": 7716.68, "duration_seconds": 12.8, "start_time": "02:08:23.880", "end_time": "02:08:36.680", "frame_count": 320, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_004_before_f192596.jpg", "start": "evidence/0006/interval_004_start_f192597.jpg", "end": "evidence/0006/interval_004_end_f192916.jpg", "after": "evidence/0006/interval_004_after_f192917.jpg"}}, {"start_frame": 196217, "end_frame": 197106, "start_seconds": 7848.68, "end_seconds": 7884.28, "duration_seconds": 35.6, "start_time": "02:10:48.680", "end_time": "02:11:24.280", "frame_count": 890, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_005_before_f196216.jpg", "start": "evidence/0006/interval_005_start_f196217.jpg", "end": "evidence/0006/interval_005_end_f197106.jpg", "after": "evidence/0006/interval_005_after_f197107.jpg"}}, {"start_frame": 198352, "end_frame": 198542, "start_seconds": 7934.08, "end_seconds": 7941.72, "duration_seconds": 7.64, "start_time": "02:12:14.080", "end_time": "02:12:21.720", "frame_count": 191, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_006_before_f198351.jpg", "start": "evidence/0006/interval_006_start_f198352.jpg", "end": "evidence/0006/interval_006_end_f198542.jpg", "after": "evidence/0006/interval_006_after_f198543.jpg"}}, {"start_frame": 204030, "end_frame": 204189, "start_seconds": 8161.2, "end_seconds": 8167.6, "duration_seconds": 6.4, "start_time": "02:16:01.200", "end_time": "02:16:07.600", "frame_count": 160, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_007_before_f204029.jpg", "start": "evidence/0006/interval_007_start_f204030.jpg", "end": "evidence/0006/interval_007_end_f204189.jpg", "after": "evidence/0006/interval_007_after_f204190.jpg"}}, {"start_frame": 222489, "end_frame": 222898, "start_seconds": 8899.56, "end_seconds": 8915.96, "duration_seconds": 16.4, "start_time": "02:28:19.560", "end_time": "02:28:35.960", "frame_count": 410, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_008_before_f222488.jpg", "start": "evidence/0006/interval_008_start_f222489.jpg", "end": "evidence/0006/interval_008_end_f222898.jpg", "after": "evidence/0006/interval_008_after_f222899.jpg"}}, {"start_frame": 239305, "end_frame": 239671, "start_seconds": 9572.2, "end_seconds": 9586.88, "duration_seconds": 14.68, "start_time": "02:39:32.200", "end_time": "02:39:46.880", "frame_count": 367, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_009_before_f239304.jpg", "start": "evidence/0006/interval_009_start_f239305.jpg", "end": "evidence/0006/interval_009_end_f239671.jpg", "after": "evidence/0006/interval_009_after_f239672.jpg"}}, {"start_frame": 241199, "end_frame": 241485, "start_seconds": 9647.96, "end_seconds": 9659.44, "duration_seconds": 11.48, "start_time": "02:40:47.960", "end_time": "02:40:59.440", "frame_count": 287, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_010_before_f241198.jpg", "start": "evidence/0006/interval_010_start_f241199.jpg", "end": "evidence/0006/interval_010_end_f241485.jpg", "after": "evidence/0006/interval_010_after_f241486.jpg"}}, {"start_frame": 247032, "end_frame": 252252, "start_seconds": 9881.28, "end_seconds": 10090.12, "duration_seconds": 208.84, "start_time": "02:44:41.280", "end_time": "02:48:10.120", "frame_count": 5221, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0006/interval_011_before_f247031.jpg", "start": "evidence/0006/interval_011_start_f247032.jpg", "end": "evidence/0006/interval_011_end_f252252.jpg", "after": "evidence/0006/interval_011_after_f252253.jpg"}}], "fps": 25.0, "frame_count": 255550, "width": 960, "height": 540, "duration_seconds": 10222.0, "decoded_frames": 255550, "coarse_samples": 10222, "coarse_candidate_count": 12, "interval_count": 11, "blue_screen": true, "blue_screen_seconds": 2802.6, "blue_screen_ratio": 0.2741733515945999, "finished_at": "2026-08-25T00:58:38.621540+08:00", "elapsed_seconds": 234.3376064300537, "video_filename": "0006 - Heico - Prokto - 7.avi"} +{"video_id": "0007", "status": "completed", "started_at": "2026-08-25T00:58:38.653328+08:00", "intervals": [{"start_frame": 2119, "end_frame": 2344, "start_seconds": 84.76, "end_seconds": 93.8, "duration_seconds": 9.04, "start_time": "00:01:24.760", "end_time": "00:01:33.800", "frame_count": 226, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_001_before_f2118.jpg", "start": "evidence/0007/interval_001_start_f2119.jpg", "end": "evidence/0007/interval_001_end_f2344.jpg", "after": "evidence/0007/interval_001_after_f2345.jpg"}}, {"start_frame": 2520, "end_frame": 2930, "start_seconds": 100.8, "end_seconds": 117.24, "duration_seconds": 16.44, "start_time": "00:01:40.800", "end_time": "00:01:57.240", "frame_count": 411, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_002_before_f2519.jpg", "start": "evidence/0007/interval_002_start_f2520.jpg", "end": "evidence/0007/interval_002_end_f2930.jpg", "after": "evidence/0007/interval_002_after_f2931.jpg"}}, {"start_frame": 77279, "end_frame": 77433, "start_seconds": 3091.16, "end_seconds": 3097.36, "duration_seconds": 6.2, "start_time": "00:51:31.160", "end_time": "00:51:37.360", "frame_count": 155, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_003_before_f77278.jpg", "start": "evidence/0007/interval_003_start_f77279.jpg", "end": "evidence/0007/interval_003_end_f77433.jpg", "after": "evidence/0007/interval_003_after_f77434.jpg"}}, {"start_frame": 125114, "end_frame": 126789, "start_seconds": 5004.56, "end_seconds": 5071.6, "duration_seconds": 67.04, "start_time": "01:23:24.560", "end_time": "01:24:31.600", "frame_count": 1676, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_004_before_f125113.jpg", "start": "evidence/0007/interval_004_start_f125114.jpg", "end": "evidence/0007/interval_004_end_f126789.jpg", "after": "evidence/0007/interval_004_after_f126790.jpg"}}, {"start_frame": 150927, "end_frame": 209117, "start_seconds": 6037.08, "end_seconds": 8364.72, "duration_seconds": 2327.64, "start_time": "01:40:37.080", "end_time": "02:19:24.720", "frame_count": 58191, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_005_before_f150926.jpg", "start": "evidence/0007/interval_005_start_f150927.jpg", "end": "evidence/0007/interval_005_end_f209117.jpg", "after": "evidence/0007/interval_005_after_f209118.jpg"}}, {"start_frame": 209372, "end_frame": 209595, "start_seconds": 8374.88, "end_seconds": 8383.84, "duration_seconds": 8.96, "start_time": "02:19:34.880", "end_time": "02:19:43.840", "frame_count": 224, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_006_before_f209371.jpg", "start": "evidence/0007/interval_006_start_f209372.jpg", "end": "evidence/0007/interval_006_end_f209595.jpg", "after": "evidence/0007/interval_006_after_f209596.jpg"}}, {"start_frame": 237291, "end_frame": 237838, "start_seconds": 9491.64, "end_seconds": 9513.56, "duration_seconds": 21.92, "start_time": "02:38:11.640", "end_time": "02:38:33.560", "frame_count": 548, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_007_before_f237290.jpg", "start": "evidence/0007/interval_007_start_f237291.jpg", "end": "evidence/0007/interval_007_end_f237838.jpg", "after": "evidence/0007/interval_007_after_f237839.jpg"}}, {"start_frame": 238028, "end_frame": 238568, "start_seconds": 9521.12, "end_seconds": 9542.76, "duration_seconds": 21.64, "start_time": "02:38:41.120", "end_time": "02:39:02.760", "frame_count": 541, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_008_before_f238027.jpg", "start": "evidence/0007/interval_008_start_f238028.jpg", "end": "evidence/0007/interval_008_end_f238568.jpg", "after": "evidence/0007/interval_008_after_f238569.jpg"}}, {"start_frame": 238736, "end_frame": 238982, "start_seconds": 9549.44, "end_seconds": 9559.32, "duration_seconds": 9.88, "start_time": "02:39:09.440", "end_time": "02:39:19.320", "frame_count": 247, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_009_before_f238735.jpg", "start": "evidence/0007/interval_009_start_f238736.jpg", "end": "evidence/0007/interval_009_end_f238982.jpg", "after": "evidence/0007/interval_009_after_f238983.jpg"}}, {"start_frame": 252643, "end_frame": 254287, "start_seconds": 10105.72, "end_seconds": 10171.52, "duration_seconds": 65.8, "start_time": "02:48:25.720", "end_time": "02:49:31.520", "frame_count": 1645, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0007/interval_010_before_f252642.jpg", "start": "evidence/0007/interval_010_start_f252643.jpg", "end": "evidence/0007/interval_010_end_f254287.jpg", "after": "evidence/0007/interval_010_after_f254288.jpg"}}], "fps": 25.0, "frame_count": 259709, "width": 960, "height": 540, "duration_seconds": 10388.36, "decoded_frames": 259709, "coarse_samples": 10389, "coarse_candidate_count": 10, "interval_count": 10, "blue_screen": true, "blue_screen_seconds": 2554.56, "blue_screen_ratio": 0.24590599478647254, "finished_at": "2026-08-25T01:02:33.387278+08:00", "elapsed_seconds": 234.73398072645068, "video_filename": "0007 - Heico - Prokto - 8.avi"} +{"video_id": "0008", "status": "completed", "started_at": "2026-08-25T01:02:33.440142+08:00", "intervals": [{"start_frame": 2459, "end_frame": 2677, "start_seconds": 98.36, "end_seconds": 107.12, "duration_seconds": 8.76, "start_time": "00:01:38.360", "end_time": "00:01:47.120", "frame_count": 219, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_001_before_f2458.jpg", "start": "evidence/0008/interval_001_start_f2459.jpg", "end": "evidence/0008/interval_001_end_f2677.jpg", "after": "evidence/0008/interval_001_after_f2678.jpg"}}, {"start_frame": 2871, "end_frame": 3130, "start_seconds": 114.84, "end_seconds": 125.24, "duration_seconds": 10.4, "start_time": "00:01:54.840", "end_time": "00:02:05.240", "frame_count": 260, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_002_before_f2870.jpg", "start": "evidence/0008/interval_002_start_f2871.jpg", "end": "evidence/0008/interval_002_end_f3130.jpg", "after": "evidence/0008/interval_002_after_f3131.jpg"}}, {"start_frame": 5009, "end_frame": 5173, "start_seconds": 200.36, "end_seconds": 206.96, "duration_seconds": 6.6, "start_time": "00:03:20.360", "end_time": "00:03:26.960", "frame_count": 165, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_003_before_f5008.jpg", "start": "evidence/0008/interval_003_start_f5009.jpg", "end": "evidence/0008/interval_003_end_f5173.jpg", "after": "evidence/0008/interval_003_after_f5174.jpg"}}, {"start_frame": 5444, "end_frame": 5697, "start_seconds": 217.76, "end_seconds": 227.92, "duration_seconds": 10.16, "start_time": "00:03:37.760", "end_time": "00:03:47.920", "frame_count": 254, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_004_before_f5443.jpg", "start": "evidence/0008/interval_004_start_f5444.jpg", "end": "evidence/0008/interval_004_end_f5697.jpg", "after": "evidence/0008/interval_004_after_f5698.jpg"}}, {"start_frame": 166621, "end_frame": 237773, "start_seconds": 6664.84, "end_seconds": 9510.96, "duration_seconds": 2846.12, "start_time": "01:51:04.840", "end_time": "02:38:30.960", "frame_count": 71153, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_005_before_f166620.jpg", "start": "evidence/0008/interval_005_start_f166621.jpg", "end": "evidence/0008/interval_005_end_f237773.jpg", "after": "evidence/0008/interval_005_after_f237774.jpg"}}, {"start_frame": 238250, "end_frame": 238399, "start_seconds": 9530.0, "end_seconds": 9536.0, "duration_seconds": 6.0, "start_time": "02:38:50.000", "end_time": "02:38:56.000", "frame_count": 150, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_006_before_f238249.jpg", "start": "evidence/0008/interval_006_start_f238250.jpg", "end": "evidence/0008/interval_006_end_f238399.jpg", "after": "evidence/0008/interval_006_after_f238400.jpg"}}, {"start_frame": 279716, "end_frame": 280051, "start_seconds": 11188.64, "end_seconds": 11202.08, "duration_seconds": 13.44, "start_time": "03:06:28.640", "end_time": "03:06:42.080", "frame_count": 336, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_007_before_f279715.jpg", "start": "evidence/0008/interval_007_start_f279716.jpg", "end": "evidence/0008/interval_007_end_f280051.jpg", "after": "evidence/0008/interval_007_after_f280052.jpg"}}, {"start_frame": 280996, "end_frame": 281540, "start_seconds": 11239.84, "end_seconds": 11261.64, "duration_seconds": 21.8, "start_time": "03:07:19.840", "end_time": "03:07:41.640", "frame_count": 545, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_008_before_f280995.jpg", "start": "evidence/0008/interval_008_start_f280996.jpg", "end": "evidence/0008/interval_008_end_f281540.jpg", "after": "evidence/0008/interval_008_after_f281541.jpg"}}, {"start_frame": 295157, "end_frame": 298333, "start_seconds": 11806.28, "end_seconds": 11933.36, "duration_seconds": 127.08, "start_time": "03:16:46.280", "end_time": "03:18:53.360", "frame_count": 3177, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_009_before_f295156.jpg", "start": "evidence/0008/interval_009_start_f295157.jpg", "end": "evidence/0008/interval_009_end_f298333.jpg", "after": "evidence/0008/interval_009_after_f298334.jpg"}}, {"start_frame": 299481, "end_frame": 299873, "start_seconds": 11979.24, "end_seconds": 11994.96, "duration_seconds": 15.72, "start_time": "03:19:39.240", "end_time": "03:19:54.960", "frame_count": 393, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_010_before_f299480.jpg", "start": "evidence/0008/interval_010_start_f299481.jpg", "end": "evidence/0008/interval_010_end_f299873.jpg", "after": "evidence/0008/interval_010_after_f299874.jpg"}}, {"start_frame": 301505, "end_frame": 302399, "start_seconds": 12060.2, "end_seconds": 12096.0, "duration_seconds": 35.8, "start_time": "03:21:00.200", "end_time": "03:21:36.000", "frame_count": 895, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_011_before_f301504.jpg", "start": "evidence/0008/interval_011_start_f301505.jpg", "end": "evidence/0008/interval_011_end_f302399.jpg", "after": "evidence/0008/interval_011_after_f302400.jpg"}}, {"start_frame": 304394, "end_frame": 305695, "start_seconds": 12175.76, "end_seconds": 12227.84, "duration_seconds": 52.08, "start_time": "03:22:55.760", "end_time": "03:23:47.840", "frame_count": 1302, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_012_before_f304393.jpg", "start": "evidence/0008/interval_012_start_f304394.jpg", "end": "evidence/0008/interval_012_end_f305695.jpg", "after": "evidence/0008/interval_012_after_f305696.jpg"}}, {"start_frame": 306196, "end_frame": 310438, "start_seconds": 12247.84, "end_seconds": 12417.56, "duration_seconds": 169.72, "start_time": "03:24:07.840", "end_time": "03:26:57.560", "frame_count": 4243, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_013_before_f306195.jpg", "start": "evidence/0008/interval_013_start_f306196.jpg", "end": "evidence/0008/interval_013_end_f310438.jpg", "after": "evidence/0008/interval_013_after_f310439.jpg"}}, {"start_frame": 313476, "end_frame": 313750, "start_seconds": 12539.04, "end_seconds": 12550.04, "duration_seconds": 11.0, "start_time": "03:28:59.040", "end_time": "03:29:10.040", "frame_count": 275, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_014_before_f313475.jpg", "start": "evidence/0008/interval_014_start_f313476.jpg", "end": "evidence/0008/interval_014_end_f313750.jpg", "after": "evidence/0008/interval_014_after_f313751.jpg"}}, {"start_frame": 315281, "end_frame": 315430, "start_seconds": 12611.24, "end_seconds": 12617.24, "duration_seconds": 6.0, "start_time": "03:30:11.240", "end_time": "03:30:17.240", "frame_count": 150, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0008/interval_015_before_f315280.jpg", "start": "evidence/0008/interval_015_start_f315281.jpg", "end": "evidence/0008/interval_015_end_f315430.jpg", "after": "evidence/0008/interval_015_after_f315431.jpg"}}], "fps": 25.0, "frame_count": 317882, "width": 960, "height": 540, "duration_seconds": 12715.28, "decoded_frames": 317882, "coarse_samples": 12716, "coarse_candidate_count": 17, "interval_count": 15, "blue_screen": true, "blue_screen_seconds": 3340.68, "blue_screen_ratio": 0.262729566317061, "finished_at": "2026-08-25T01:07:59.992614+08:00", "elapsed_seconds": 326.55251191556454, "video_filename": "0008 - Heico - Prokto - 9.avi"} +{"video_id": "0009", "status": "completed", "started_at": "2026-08-25T01:08:00.109833+08:00", "intervals": [{"start_frame": 7060, "end_frame": 8007, "start_seconds": 282.4, "end_seconds": 320.32, "duration_seconds": 37.92, "start_time": "00:04:42.400", "end_time": "00:05:20.320", "frame_count": 948, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_001_before_f7059.jpg", "start": "evidence/0009/interval_001_start_f7060.jpg", "end": "evidence/0009/interval_001_end_f8007.jpg", "after": "evidence/0009/interval_001_after_f8008.jpg"}}, {"start_frame": 8244, "end_frame": 8518, "start_seconds": 329.76, "end_seconds": 340.76, "duration_seconds": 11.0, "start_time": "00:05:29.760", "end_time": "00:05:40.760", "frame_count": 275, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_002_before_f8243.jpg", "start": "evidence/0009/interval_002_start_f8244.jpg", "end": "evidence/0009/interval_002_end_f8518.jpg", "after": "evidence/0009/interval_002_after_f8519.jpg"}}, {"start_frame": 10905, "end_frame": 11723, "start_seconds": 436.2, "end_seconds": 468.96, "duration_seconds": 32.76, "start_time": "00:07:16.200", "end_time": "00:07:48.960", "frame_count": 819, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_003_before_f10904.jpg", "start": "evidence/0009/interval_003_start_f10905.jpg", "end": "evidence/0009/interval_003_end_f11723.jpg", "after": "evidence/0009/interval_003_after_f11724.jpg"}}, {"start_frame": 12679, "end_frame": 12928, "start_seconds": 507.16, "end_seconds": 517.16, "duration_seconds": 10.0, "start_time": "00:08:27.160", "end_time": "00:08:37.160", "frame_count": 250, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_004_before_f12678.jpg", "start": "evidence/0009/interval_004_start_f12679.jpg", "end": "evidence/0009/interval_004_end_f12928.jpg", "after": "evidence/0009/interval_004_after_f12929.jpg"}}, {"start_frame": 13404, "end_frame": 13908, "start_seconds": 536.16, "end_seconds": 556.36, "duration_seconds": 20.2, "start_time": "00:08:56.160", "end_time": "00:09:16.360", "frame_count": 505, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_005_before_f13403.jpg", "start": "evidence/0009/interval_005_start_f13404.jpg", "end": "evidence/0009/interval_005_end_f13908.jpg", "after": "evidence/0009/interval_005_after_f13909.jpg"}}, {"start_frame": 53068, "end_frame": 58670, "start_seconds": 2122.72, "end_seconds": 2346.84, "duration_seconds": 224.12, "start_time": "00:35:22.720", "end_time": "00:39:06.840", "frame_count": 5603, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_006_before_f53067.jpg", "start": "evidence/0009/interval_006_start_f53068.jpg", "end": "evidence/0009/interval_006_end_f58670.jpg", "after": "evidence/0009/interval_006_after_f58671.jpg"}}, {"start_frame": 63976, "end_frame": 65220, "start_seconds": 2559.04, "end_seconds": 2608.84, "duration_seconds": 49.8, "start_time": "00:42:39.040", "end_time": "00:43:28.840", "frame_count": 1245, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_007_before_f63975.jpg", "start": "evidence/0009/interval_007_start_f63976.jpg", "end": "evidence/0009/interval_007_end_f65220.jpg", "after": "evidence/0009/interval_007_after_f65221.jpg"}}, {"start_frame": 65405, "end_frame": 65853, "start_seconds": 2616.2, "end_seconds": 2634.16, "duration_seconds": 17.96, "start_time": "00:43:36.200", "end_time": "00:43:54.160", "frame_count": 449, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_008_before_f65404.jpg", "start": "evidence/0009/interval_008_start_f65405.jpg", "end": "evidence/0009/interval_008_end_f65853.jpg", "after": "evidence/0009/interval_008_after_f65854.jpg"}}, {"start_frame": 66390, "end_frame": 67030, "start_seconds": 2655.6, "end_seconds": 2681.24, "duration_seconds": 25.64, "start_time": "00:44:15.600", "end_time": "00:44:41.240", "frame_count": 641, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_009_before_f66389.jpg", "start": "evidence/0009/interval_009_start_f66390.jpg", "end": "evidence/0009/interval_009_end_f67030.jpg", "after": "evidence/0009/interval_009_after_f67031.jpg"}}, {"start_frame": 85249, "end_frame": 86243, "start_seconds": 3409.96, "end_seconds": 3449.76, "duration_seconds": 39.8, "start_time": "00:56:49.960", "end_time": "00:57:29.760", "frame_count": 995, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_010_before_f85248.jpg", "start": "evidence/0009/interval_010_start_f85249.jpg", "end": "evidence/0009/interval_010_end_f86243.jpg", "after": "evidence/0009/interval_010_after_f86244.jpg"}}, {"start_frame": 91754, "end_frame": 92359, "start_seconds": 3670.16, "end_seconds": 3694.4, "duration_seconds": 24.24, "start_time": "01:01:10.160", "end_time": "01:01:34.400", "frame_count": 606, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_011_before_f91753.jpg", "start": "evidence/0009/interval_011_start_f91754.jpg", "end": "evidence/0009/interval_011_end_f92359.jpg", "after": "evidence/0009/interval_011_after_f92360.jpg"}}, {"start_frame": 92636, "end_frame": 93048, "start_seconds": 3705.44, "end_seconds": 3721.96, "duration_seconds": 16.52, "start_time": "01:01:45.440", "end_time": "01:02:01.960", "frame_count": 413, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_012_before_f92635.jpg", "start": "evidence/0009/interval_012_start_f92636.jpg", "end": "evidence/0009/interval_012_end_f93048.jpg", "after": "evidence/0009/interval_012_after_f93049.jpg"}}, {"start_frame": 94012, "end_frame": 94245, "start_seconds": 3760.48, "end_seconds": 3769.84, "duration_seconds": 9.36, "start_time": "01:02:40.480", "end_time": "01:02:49.840", "frame_count": 234, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_013_before_f94011.jpg", "start": "evidence/0009/interval_013_start_f94012.jpg", "end": "evidence/0009/interval_013_end_f94245.jpg", "after": "evidence/0009/interval_013_after_f94246.jpg"}}, {"start_frame": 125101, "end_frame": 126017, "start_seconds": 5004.04, "end_seconds": 5040.72, "duration_seconds": 36.68, "start_time": "01:23:24.040", "end_time": "01:24:00.720", "frame_count": 917, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_014_before_f125100.jpg", "start": "evidence/0009/interval_014_start_f125101.jpg", "end": "evidence/0009/interval_014_end_f126017.jpg", "after": "evidence/0009/interval_014_after_f126018.jpg"}}, {"start_frame": 126104, "end_frame": 126615, "start_seconds": 5044.16, "end_seconds": 5064.64, "duration_seconds": 20.48, "start_time": "01:24:04.160", "end_time": "01:24:24.640", "frame_count": 512, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_015_before_f126103.jpg", "start": "evidence/0009/interval_015_start_f126104.jpg", "end": "evidence/0009/interval_015_end_f126615.jpg", "after": "evidence/0009/interval_015_after_f126616.jpg"}}, {"start_frame": 145892, "end_frame": 146607, "start_seconds": 5835.68, "end_seconds": 5864.32, "duration_seconds": 28.64, "start_time": "01:37:15.680", "end_time": "01:37:44.320", "frame_count": 716, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_016_before_f145891.jpg", "start": "evidence/0009/interval_016_start_f145892.jpg", "end": "evidence/0009/interval_016_end_f146607.jpg", "after": "evidence/0009/interval_016_after_f146608.jpg"}}, {"start_frame": 147498, "end_frame": 148638, "start_seconds": 5899.92, "end_seconds": 5945.56, "duration_seconds": 45.64, "start_time": "01:38:19.920", "end_time": "01:39:05.560", "frame_count": 1141, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_017_before_f147497.jpg", "start": "evidence/0009/interval_017_start_f147498.jpg", "end": "evidence/0009/interval_017_end_f148638.jpg", "after": "evidence/0009/interval_017_after_f148639.jpg"}}, {"start_frame": 151144, "end_frame": 153672, "start_seconds": 6045.76, "end_seconds": 6146.92, "duration_seconds": 101.16, "start_time": "01:40:45.760", "end_time": "01:42:26.920", "frame_count": 2529, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_018_before_f151143.jpg", "start": "evidence/0009/interval_018_start_f151144.jpg", "end": "evidence/0009/interval_018_end_f153672.jpg", "after": "evidence/0009/interval_018_after_f153673.jpg"}}, {"start_frame": 153819, "end_frame": 154897, "start_seconds": 6152.76, "end_seconds": 6195.92, "duration_seconds": 43.16, "start_time": "01:42:32.760", "end_time": "01:43:15.920", "frame_count": 1079, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_019_before_f153818.jpg", "start": "evidence/0009/interval_019_start_f153819.jpg", "end": "evidence/0009/interval_019_end_f154897.jpg", "after": "evidence/0009/interval_019_after_f154898.jpg"}}, {"start_frame": 156864, "end_frame": 157369, "start_seconds": 6274.56, "end_seconds": 6294.8, "duration_seconds": 20.24, "start_time": "01:44:34.560", "end_time": "01:44:54.800", "frame_count": 506, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_020_before_f156863.jpg", "start": "evidence/0009/interval_020_start_f156864.jpg", "end": "evidence/0009/interval_020_end_f157369.jpg", "after": "evidence/0009/interval_020_after_f157370.jpg"}}, {"start_frame": 158007, "end_frame": 158420, "start_seconds": 6320.28, "end_seconds": 6336.84, "duration_seconds": 16.56, "start_time": "01:45:20.280", "end_time": "01:45:36.840", "frame_count": 414, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_021_before_f158006.jpg", "start": "evidence/0009/interval_021_start_f158007.jpg", "end": "evidence/0009/interval_021_end_f158420.jpg", "after": "evidence/0009/interval_021_after_f158421.jpg"}}, {"start_frame": 171517, "end_frame": 171713, "start_seconds": 6860.68, "end_seconds": 6868.56, "duration_seconds": 7.88, "start_time": "01:54:20.680", "end_time": "01:54:28.560", "frame_count": 197, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_022_before_f171516.jpg", "start": "evidence/0009/interval_022_start_f171517.jpg", "end": "evidence/0009/interval_022_end_f171713.jpg", "after": "evidence/0009/interval_022_after_f171714.jpg"}}, {"start_frame": 185776, "end_frame": 187204, "start_seconds": 7431.04, "end_seconds": 7488.2, "duration_seconds": 57.16, "start_time": "02:03:51.040", "end_time": "02:04:48.200", "frame_count": 1429, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_023_before_f185775.jpg", "start": "evidence/0009/interval_023_start_f185776.jpg", "end": "evidence/0009/interval_023_end_f187204.jpg", "after": "evidence/0009/interval_023_after_f187205.jpg"}}, {"start_frame": 187469, "end_frame": 187733, "start_seconds": 7498.76, "end_seconds": 7509.36, "duration_seconds": 10.6, "start_time": "02:04:58.760", "end_time": "02:05:09.360", "frame_count": 265, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_024_before_f187468.jpg", "start": "evidence/0009/interval_024_start_f187469.jpg", "end": "evidence/0009/interval_024_end_f187733.jpg", "after": "evidence/0009/interval_024_after_f187734.jpg"}}, {"start_frame": 206430, "end_frame": 207747, "start_seconds": 8257.2, "end_seconds": 8309.92, "duration_seconds": 52.72, "start_time": "02:17:37.200", "end_time": "02:18:29.920", "frame_count": 1318, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_025_before_f206429.jpg", "start": "evidence/0009/interval_025_start_f206430.jpg", "end": "evidence/0009/interval_025_end_f207747.jpg", "after": "evidence/0009/interval_025_after_f207748.jpg"}}, {"start_frame": 207944, "end_frame": 208182, "start_seconds": 8317.76, "end_seconds": 8327.32, "duration_seconds": 9.56, "start_time": "02:18:37.760", "end_time": "02:18:47.320", "frame_count": 239, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_026_before_f207943.jpg", "start": "evidence/0009/interval_026_start_f207944.jpg", "end": "evidence/0009/interval_026_end_f208182.jpg", "after": "evidence/0009/interval_026_after_f208183.jpg"}}, {"start_frame": 214536, "end_frame": 214878, "start_seconds": 8581.44, "end_seconds": 8595.16, "duration_seconds": 13.72, "start_time": "02:23:01.440", "end_time": "02:23:15.160", "frame_count": 343, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_027_before_f214535.jpg", "start": "evidence/0009/interval_027_start_f214536.jpg", "end": "evidence/0009/interval_027_end_f214878.jpg", "after": "evidence/0009/interval_027_after_f214879.jpg"}}, {"start_frame": 214972, "end_frame": 215438, "start_seconds": 8598.88, "end_seconds": 8617.56, "duration_seconds": 18.68, "start_time": "02:23:18.880", "end_time": "02:23:37.560", "frame_count": 467, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_028_before_f214971.jpg", "start": "evidence/0009/interval_028_start_f214972.jpg", "end": "evidence/0009/interval_028_end_f215438.jpg", "after": "evidence/0009/interval_028_after_f215439.jpg"}}, {"start_frame": 216309, "end_frame": 216908, "start_seconds": 8652.36, "end_seconds": 8676.36, "duration_seconds": 24.0, "start_time": "02:24:12.360", "end_time": "02:24:36.360", "frame_count": 600, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_029_before_f216308.jpg", "start": "evidence/0009/interval_029_start_f216309.jpg", "end": "evidence/0009/interval_029_end_f216908.jpg", "after": "evidence/0009/interval_029_after_f216909.jpg"}}, {"start_frame": 219259, "end_frame": 219518, "start_seconds": 8770.36, "end_seconds": 8780.76, "duration_seconds": 10.4, "start_time": "02:26:10.360", "end_time": "02:26:20.760", "frame_count": 260, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_030_before_f219258.jpg", "start": "evidence/0009/interval_030_start_f219259.jpg", "end": "evidence/0009/interval_030_end_f219518.jpg", "after": "evidence/0009/interval_030_after_f219519.jpg"}}, {"start_frame": 223627, "end_frame": 223803, "start_seconds": 8945.08, "end_seconds": 8952.16, "duration_seconds": 7.08, "start_time": "02:29:05.080", "end_time": "02:29:12.160", "frame_count": 177, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_031_before_f223626.jpg", "start": "evidence/0009/interval_031_start_f223627.jpg", "end": "evidence/0009/interval_031_end_f223803.jpg", "after": "evidence/0009/interval_031_after_f223804.jpg"}}, {"start_frame": 229414, "end_frame": 229663, "start_seconds": 9176.56, "end_seconds": 9186.56, "duration_seconds": 10.0, "start_time": "02:32:56.560", "end_time": "02:33:06.560", "frame_count": 250, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_032_before_f229413.jpg", "start": "evidence/0009/interval_032_start_f229414.jpg", "end": "evidence/0009/interval_032_end_f229663.jpg", "after": "evidence/0009/interval_032_after_f229664.jpg"}}, {"start_frame": 280676, "end_frame": 283232, "start_seconds": 11227.04, "end_seconds": 11329.32, "duration_seconds": 102.28, "start_time": "03:07:07.040", "end_time": "03:08:49.320", "frame_count": 2557, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_033_before_f280675.jpg", "start": "evidence/0009/interval_033_start_f280676.jpg", "end": "evidence/0009/interval_033_end_f283232.jpg", "after": "evidence/0009/interval_033_after_f283233.jpg"}}, {"start_frame": 308970, "end_frame": 373914, "start_seconds": 12358.8, "end_seconds": 14956.6, "duration_seconds": 2597.8, "start_time": "03:25:58.800", "end_time": "04:09:16.600", "frame_count": 64945, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_034_before_f308969.jpg", "start": "evidence/0009/interval_034_start_f308970.jpg", "end": "evidence/0009/interval_034_end_f373914.jpg", "after": "evidence/0009/interval_034_after_f373915.jpg"}}, {"start_frame": 374164, "end_frame": 374543, "start_seconds": 14966.56, "end_seconds": 14981.76, "duration_seconds": 15.2, "start_time": "04:09:26.560", "end_time": "04:09:41.760", "frame_count": 380, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_035_before_f374163.jpg", "start": "evidence/0009/interval_035_start_f374164.jpg", "end": "evidence/0009/interval_035_end_f374543.jpg", "after": "evidence/0009/interval_035_after_f374544.jpg"}}, {"start_frame": 387421, "end_frame": 388107, "start_seconds": 15496.84, "end_seconds": 15524.32, "duration_seconds": 27.48, "start_time": "04:18:16.840", "end_time": "04:18:44.320", "frame_count": 687, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_036_before_f387420.jpg", "start": "evidence/0009/interval_036_start_f387421.jpg", "end": "evidence/0009/interval_036_end_f388107.jpg", "after": "evidence/0009/interval_036_after_f388108.jpg"}}, {"start_frame": 404192, "end_frame": 404443, "start_seconds": 16167.68, "end_seconds": 16177.76, "duration_seconds": 10.08, "start_time": "04:29:27.680", "end_time": "04:29:37.760", "frame_count": 252, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_037_before_f404191.jpg", "start": "evidence/0009/interval_037_start_f404192.jpg", "end": "evidence/0009/interval_037_end_f404443.jpg", "after": "evidence/0009/interval_037_after_f404444.jpg"}}, {"start_frame": 407105, "end_frame": 407310, "start_seconds": 16284.2, "end_seconds": 16292.44, "duration_seconds": 8.24, "start_time": "04:31:24.200", "end_time": "04:31:32.440", "frame_count": 206, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_038_before_f407104.jpg", "start": "evidence/0009/interval_038_start_f407105.jpg", "end": "evidence/0009/interval_038_end_f407310.jpg", "after": "evidence/0009/interval_038_after_f407311.jpg"}}, {"start_frame": 432003, "end_frame": 441298, "start_seconds": 17280.12, "end_seconds": 17651.96, "duration_seconds": 371.84, "start_time": "04:48:00.120", "end_time": "04:54:11.960", "frame_count": 9296, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_039_before_f432002.jpg", "start": "evidence/0009/interval_039_start_f432003.jpg", "end": "evidence/0009/interval_039_end_f441298.jpg", "after": "evidence/0009/interval_039_after_f441299.jpg"}}, {"start_frame": 442866, "end_frame": 443022, "start_seconds": 17714.64, "end_seconds": 17720.92, "duration_seconds": 6.28, "start_time": "04:55:14.640", "end_time": "04:55:20.920", "frame_count": 157, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0009/interval_040_before_f442865.jpg", "start": "evidence/0009/interval_040_start_f442866.jpg", "end": "evidence/0009/interval_040_end_f443022.jpg", "after": "evidence/0009/interval_040_after_f443023.jpg"}}], "fps": 25.0, "frame_count": 444515, "width": 960, "height": 540, "duration_seconds": 17780.6, "decoded_frames": 444515, "coarse_samples": 17781, "coarse_candidate_count": 41, "interval_count": 40, "blue_screen": true, "blue_screen_seconds": 4192.879999999999, "blue_screen_ratio": 0.23581206483470746, "finished_at": "2026-08-25T01:15:52.834328+08:00", "elapsed_seconds": 472.72453361563385, "video_filename": "0009 - Heico - Prokto - 10.avi"} +{"video_id": "0010", "status": "completed", "started_at": "2026-08-25T01:15:52.902143+08:00", "intervals": [{"start_frame": 41235, "end_frame": 41710, "start_seconds": 1649.4, "end_seconds": 1668.44, "duration_seconds": 19.04, "start_time": "00:27:29.400", "end_time": "00:27:48.440", "frame_count": 476, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_001_before_f41234.jpg", "start": "evidence/0010/interval_001_start_f41235.jpg", "end": "evidence/0010/interval_001_end_f41710.jpg", "after": "evidence/0010/interval_001_after_f41711.jpg"}}, {"start_frame": 49196, "end_frame": 50101, "start_seconds": 1967.84, "end_seconds": 2004.08, "duration_seconds": 36.24, "start_time": "00:32:47.840", "end_time": "00:33:24.080", "frame_count": 906, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_002_before_f49195.jpg", "start": "evidence/0010/interval_002_start_f49196.jpg", "end": "evidence/0010/interval_002_end_f50101.jpg", "after": "evidence/0010/interval_002_after_f50102.jpg"}}, {"start_frame": 157223, "end_frame": 158078, "start_seconds": 6288.92, "end_seconds": 6323.16, "duration_seconds": 34.24, "start_time": "01:44:48.920", "end_time": "01:45:23.160", "frame_count": 856, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_003_before_f157222.jpg", "start": "evidence/0010/interval_003_start_f157223.jpg", "end": "evidence/0010/interval_003_end_f158078.jpg", "after": "evidence/0010/interval_003_after_f158079.jpg"}}, {"start_frame": 160177, "end_frame": 161891, "start_seconds": 6407.08, "end_seconds": 6475.68, "duration_seconds": 68.6, "start_time": "01:46:47.080", "end_time": "01:47:55.680", "frame_count": 1715, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_004_before_f160176.jpg", "start": "evidence/0010/interval_004_start_f160177.jpg", "end": "evidence/0010/interval_004_end_f161891.jpg", "after": "evidence/0010/interval_004_after_f161892.jpg"}}, {"start_frame": 165333, "end_frame": 165802, "start_seconds": 6613.32, "end_seconds": 6632.12, "duration_seconds": 18.8, "start_time": "01:50:13.320", "end_time": "01:50:32.120", "frame_count": 470, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_005_before_f165332.jpg", "start": "evidence/0010/interval_005_start_f165333.jpg", "end": "evidence/0010/interval_005_end_f165802.jpg", "after": "evidence/0010/interval_005_after_f165803.jpg"}}, {"start_frame": 265363, "end_frame": 265733, "start_seconds": 10614.52, "end_seconds": 10629.36, "duration_seconds": 14.84, "start_time": "02:56:54.520", "end_time": "02:57:09.360", "frame_count": 371, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_006_before_f265362.jpg", "start": "evidence/0010/interval_006_start_f265363.jpg", "end": "evidence/0010/interval_006_end_f265733.jpg", "after": "evidence/0010/interval_006_after_f265734.jpg"}}, {"start_frame": 265907, "end_frame": 336405, "start_seconds": 10636.28, "end_seconds": 13456.24, "duration_seconds": 2819.96, "start_time": "02:57:16.280", "end_time": "03:44:16.240", "frame_count": 70499, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_007_before_f265906.jpg", "start": "evidence/0010/interval_007_start_f265907.jpg", "end": "evidence/0010/interval_007_end_f336405.jpg", "after": "evidence/0010/interval_007_after_f336406.jpg"}}, {"start_frame": 354902, "end_frame": 355116, "start_seconds": 14196.08, "end_seconds": 14204.68, "duration_seconds": 8.6, "start_time": "03:56:36.080", "end_time": "03:56:44.680", "frame_count": 215, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_008_before_f354901.jpg", "start": "evidence/0010/interval_008_start_f354902.jpg", "end": "evidence/0010/interval_008_end_f355116.jpg", "after": "evidence/0010/interval_008_after_f355117.jpg"}}, {"start_frame": 371415, "end_frame": 371767, "start_seconds": 14856.6, "end_seconds": 14870.72, "duration_seconds": 14.12, "start_time": "04:07:36.600", "end_time": "04:07:50.720", "frame_count": 353, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_009_before_f371414.jpg", "start": "evidence/0010/interval_009_start_f371415.jpg", "end": "evidence/0010/interval_009_end_f371767.jpg", "after": "evidence/0010/interval_009_after_f371768.jpg"}}, {"start_frame": 373541, "end_frame": 373981, "start_seconds": 14941.64, "end_seconds": 14959.28, "duration_seconds": 17.64, "start_time": "04:09:01.640", "end_time": "04:09:19.280", "frame_count": 441, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_010_before_f373540.jpg", "start": "evidence/0010/interval_010_start_f373541.jpg", "end": "evidence/0010/interval_010_end_f373981.jpg", "after": "evidence/0010/interval_010_after_f373982.jpg"}}, {"start_frame": 376427, "end_frame": 376726, "start_seconds": 15057.08, "end_seconds": 15069.08, "duration_seconds": 12.0, "start_time": "04:10:57.080", "end_time": "04:11:09.080", "frame_count": 300, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_011_before_f376426.jpg", "start": "evidence/0010/interval_011_start_f376427.jpg", "end": "evidence/0010/interval_011_end_f376726.jpg", "after": "evidence/0010/interval_011_after_f376727.jpg"}}, {"start_frame": 397827, "end_frame": 398219, "start_seconds": 15913.08, "end_seconds": 15928.8, "duration_seconds": 15.72, "start_time": "04:25:13.080", "end_time": "04:25:28.800", "frame_count": 393, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_012_before_f397826.jpg", "start": "evidence/0010/interval_012_start_f397827.jpg", "end": "evidence/0010/interval_012_end_f398219.jpg", "after": "evidence/0010/interval_012_after_f398220.jpg"}}, {"start_frame": 421164, "end_frame": 429147, "start_seconds": 16846.56, "end_seconds": 17165.92, "duration_seconds": 319.36, "start_time": "04:40:46.560", "end_time": "04:46:05.920", "frame_count": 7984, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0010/interval_013_before_f421163.jpg", "start": "evidence/0010/interval_013_start_f421164.jpg", "end": "evidence/0010/interval_013_end_f429147.jpg", "after": "evidence/0010/interval_013_after_f429148.jpg"}}], "fps": 25.0, "frame_count": 430151, "width": 960, "height": 540, "duration_seconds": 17206.04, "decoded_frames": 430151, "coarse_samples": 17207, "coarse_candidate_count": 13, "interval_count": 13, "blue_screen": true, "blue_screen_seconds": 3399.16, "blue_screen_ratio": 0.1975562070063768, "finished_at": "2026-08-25T01:23:02.496188+08:00", "elapsed_seconds": 429.59408420324326, "video_filename": "0010 - Heico - Rektum - 1.avi"} +{"video_id": "0011", "status": "completed", "started_at": "2026-08-25T01:23:02.511538+08:00", "intervals": [{"start_frame": 7043, "end_frame": 7277, "start_seconds": 281.72, "end_seconds": 291.12, "duration_seconds": 9.4, "start_time": "00:04:41.720", "end_time": "00:04:51.120", "frame_count": 235, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_001_before_f7042.jpg", "start": "evidence/0011/interval_001_start_f7043.jpg", "end": "evidence/0011/interval_001_end_f7277.jpg", "after": "evidence/0011/interval_001_after_f7278.jpg"}}, {"start_frame": 176925, "end_frame": 177260, "start_seconds": 7077.0, "end_seconds": 7090.44, "duration_seconds": 13.44, "start_time": "01:57:57.000", "end_time": "01:58:10.440", "frame_count": 336, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_002_before_f176924.jpg", "start": "evidence/0011/interval_002_start_f176925.jpg", "end": "evidence/0011/interval_002_end_f177260.jpg", "after": "evidence/0011/interval_002_after_f177261.jpg"}}, {"start_frame": 183595, "end_frame": 232058, "start_seconds": 7343.8, "end_seconds": 9282.36, "duration_seconds": 1938.56, "start_time": "02:02:23.800", "end_time": "02:34:42.360", "frame_count": 48464, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_003_before_f183594.jpg", "start": "evidence/0011/interval_003_start_f183595.jpg", "end": "evidence/0011/interval_003_end_f232058.jpg", "after": "evidence/0011/interval_003_after_f232059.jpg"}}, {"start_frame": 239127, "end_frame": 240057, "start_seconds": 9565.08, "end_seconds": 9602.32, "duration_seconds": 37.24, "start_time": "02:39:25.080", "end_time": "02:40:02.320", "frame_count": 931, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_004_before_f239126.jpg", "start": "evidence/0011/interval_004_start_f239127.jpg", "end": "evidence/0011/interval_004_end_f240057.jpg", "after": "evidence/0011/interval_004_after_f240058.jpg"}}, {"start_frame": 240627, "end_frame": 241173, "start_seconds": 9625.08, "end_seconds": 9646.96, "duration_seconds": 21.88, "start_time": "02:40:25.080", "end_time": "02:40:46.960", "frame_count": 547, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_005_before_f240626.jpg", "start": "evidence/0011/interval_005_start_f240627.jpg", "end": "evidence/0011/interval_005_end_f241173.jpg", "after": "evidence/0011/interval_005_after_f241174.jpg"}}, {"start_frame": 247248, "end_frame": 247415, "start_seconds": 9889.92, "end_seconds": 9896.64, "duration_seconds": 6.72, "start_time": "02:44:49.920", "end_time": "02:44:56.640", "frame_count": 168, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_006_before_f247247.jpg", "start": "evidence/0011/interval_006_start_f247248.jpg", "end": "evidence/0011/interval_006_end_f247415.jpg", "after": "evidence/0011/interval_006_after_f247416.jpg"}}, {"start_frame": 251653, "end_frame": 254734, "start_seconds": 10066.12, "end_seconds": 10189.4, "duration_seconds": 123.28, "start_time": "02:47:46.120", "end_time": "02:49:49.400", "frame_count": 3082, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_007_before_f251652.jpg", "start": "evidence/0011/interval_007_start_f251653.jpg", "end": "evidence/0011/interval_007_end_f254734.jpg", "after": "evidence/0011/interval_007_after_f254735.jpg"}}, {"start_frame": 256223, "end_frame": 256422, "start_seconds": 10248.92, "end_seconds": 10256.92, "duration_seconds": 8.0, "start_time": "02:50:48.920", "end_time": "02:50:56.920", "frame_count": 200, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_008_before_f256222.jpg", "start": "evidence/0011/interval_008_start_f256223.jpg", "end": "evidence/0011/interval_008_end_f256422.jpg", "after": "evidence/0011/interval_008_after_f256423.jpg"}}, {"start_frame": 266223, "end_frame": 266459, "start_seconds": 10648.92, "end_seconds": 10658.4, "duration_seconds": 9.48, "start_time": "02:57:28.920", "end_time": "02:57:38.400", "frame_count": 237, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_009_before_f266222.jpg", "start": "evidence/0011/interval_009_start_f266223.jpg", "end": "evidence/0011/interval_009_end_f266459.jpg", "after": "evidence/0011/interval_009_after_f266460.jpg"}}, {"start_frame": 309058, "end_frame": 309278, "start_seconds": 12362.32, "end_seconds": 12371.16, "duration_seconds": 8.84, "start_time": "03:26:02.320", "end_time": "03:26:11.160", "frame_count": 221, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_010_before_f309057.jpg", "start": "evidence/0011/interval_010_start_f309058.jpg", "end": "evidence/0011/interval_010_end_f309278.jpg", "after": "evidence/0011/interval_010_after_f309279.jpg"}}, {"start_frame": 330794, "end_frame": 331069, "start_seconds": 13231.76, "end_seconds": 13242.8, "duration_seconds": 11.04, "start_time": "03:40:31.760", "end_time": "03:40:42.800", "frame_count": 276, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0011/interval_011_before_f330793.jpg", "start": "evidence/0011/interval_011_start_f330794.jpg", "end": "evidence/0011/interval_011_end_f331069.jpg", "after": "evidence/0011/interval_011_after_f331070.jpg"}}], "fps": 25.0, "frame_count": 344761, "width": 960, "height": 540, "duration_seconds": 13790.44, "decoded_frames": 344761, "coarse_samples": 13791, "coarse_candidate_count": 12, "interval_count": 11, "blue_screen": true, "blue_screen_seconds": 2187.88, "blue_screen_ratio": 0.1586519356887815, "finished_at": "2026-08-25T01:28:16.004720+08:00", "elapsed_seconds": 313.49321262538433, "video_filename": "0011 - Heico - Rektum - 2.avi"} +{"video_id": "0012", "status": "completed", "started_at": "2026-08-25T01:28:16.024554+08:00", "intervals": [{"start_frame": 18110, "end_frame": 18492, "start_seconds": 724.4, "end_seconds": 739.72, "duration_seconds": 15.32, "start_time": "00:12:04.400", "end_time": "00:12:19.720", "frame_count": 383, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0012/interval_001_before_f18109.jpg", "start": "evidence/0012/interval_001_start_f18110.jpg", "end": "evidence/0012/interval_001_end_f18492.jpg", "after": "evidence/0012/interval_001_after_f18493.jpg"}}, {"start_frame": 150983, "end_frame": 191811, "start_seconds": 6039.32, "end_seconds": 7672.48, "duration_seconds": 1633.16, "start_time": "01:40:39.320", "end_time": "02:07:52.480", "frame_count": 40829, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0012/interval_002_before_f150982.jpg", "start": "evidence/0012/interval_002_start_f150983.jpg", "end": "evidence/0012/interval_002_end_f191811.jpg", "after": "evidence/0012/interval_002_after_f191812.jpg"}}, {"start_frame": 211660, "end_frame": 211895, "start_seconds": 8466.4, "end_seconds": 8475.84, "duration_seconds": 9.44, "start_time": "02:21:06.400", "end_time": "02:21:15.840", "frame_count": 236, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0012/interval_003_before_f211659.jpg", "start": "evidence/0012/interval_003_start_f211660.jpg", "end": "evidence/0012/interval_003_end_f211895.jpg", "after": "evidence/0012/interval_003_after_f211896.jpg"}}], "fps": 25.0, "frame_count": 287950, "width": 960, "height": 540, "duration_seconds": 11518.0, "decoded_frames": 287950, "coarse_samples": 11518, "coarse_candidate_count": 4, "interval_count": 3, "blue_screen": true, "blue_screen_seconds": 1657.92, "blue_screen_ratio": 0.14394165653759333, "finished_at": "2026-08-25T01:32:10.505557+08:00", "elapsed_seconds": 234.48103972896934, "video_filename": "0012 - Heico - Rektum - 3.avi"} +{"video_id": "0013", "status": "completed", "started_at": "2026-08-25T01:32:10.525373+08:00", "intervals": [{"start_frame": 8189, "end_frame": 8713, "start_seconds": 327.56, "end_seconds": 348.56, "duration_seconds": 21.0, "start_time": "00:05:27.560", "end_time": "00:05:48.560", "frame_count": 525, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0013/interval_001_before_f8188.jpg", "start": "evidence/0013/interval_001_start_f8189.jpg", "end": "evidence/0013/interval_001_end_f8713.jpg", "after": "evidence/0013/interval_001_after_f8714.jpg"}}, {"start_frame": 9174, "end_frame": 9696, "start_seconds": 366.96, "end_seconds": 387.88, "duration_seconds": 20.92, "start_time": "00:06:06.960", "end_time": "00:06:27.880", "frame_count": 523, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0013/interval_002_before_f9173.jpg", "start": "evidence/0013/interval_002_start_f9174.jpg", "end": "evidence/0013/interval_002_end_f9696.jpg", "after": "evidence/0013/interval_002_after_f9697.jpg"}}, {"start_frame": 12253, "end_frame": 13148, "start_seconds": 490.12, "end_seconds": 525.96, "duration_seconds": 35.84, "start_time": "00:08:10.120", "end_time": "00:08:45.960", "frame_count": 896, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0013/interval_003_before_f12252.jpg", "start": "evidence/0013/interval_003_start_f12253.jpg", "end": "evidence/0013/interval_003_end_f13148.jpg", "after": "evidence/0013/interval_003_after_f13149.jpg"}}, {"start_frame": 16396, "end_frame": 16578, "start_seconds": 655.84, "end_seconds": 663.16, "duration_seconds": 7.32, "start_time": "00:10:55.840", "end_time": "00:11:03.160", "frame_count": 183, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0013/interval_004_before_f16395.jpg", "start": "evidence/0013/interval_004_start_f16396.jpg", "end": "evidence/0013/interval_004_end_f16578.jpg", "after": "evidence/0013/interval_004_after_f16579.jpg"}}, {"start_frame": 164401, "end_frame": 164911, "start_seconds": 6576.04, "end_seconds": 6596.48, "duration_seconds": 20.44, "start_time": "01:49:36.040", "end_time": "01:49:56.480", "frame_count": 511, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0013/interval_005_before_f164400.jpg", "start": "evidence/0013/interval_005_start_f164401.jpg", "end": "evidence/0013/interval_005_end_f164911.jpg", "after": "evidence/0013/interval_005_after_f164912.jpg"}}, {"start_frame": 168791, "end_frame": 208270, "start_seconds": 6751.64, "end_seconds": 8330.84, "duration_seconds": 1579.2, "start_time": "01:52:31.640", "end_time": "02:18:50.840", "frame_count": 39480, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0013/interval_006_before_f168790.jpg", "start": "evidence/0013/interval_006_start_f168791.jpg", "end": "evidence/0013/interval_006_end_f208270.jpg", "after": "evidence/0013/interval_006_after_f208271.jpg"}}], "fps": 25.0, "frame_count": 223365, "width": 960, "height": 540, "duration_seconds": 8934.6, "decoded_frames": 223365, "coarse_samples": 8935, "coarse_candidate_count": 6, "interval_count": 6, "blue_screen": true, "blue_screen_seconds": 1684.72, "blue_screen_ratio": 0.18856132339444406, "finished_at": "2026-08-25T01:35:46.917137+08:00", "elapsed_seconds": 216.3917945791036, "video_filename": "0013 - Heico - Rektum - 4.avi"} +{"video_id": "0014", "status": "completed", "started_at": "2026-08-25T01:35:46.941847+08:00", "intervals": [{"start_frame": 1061, "end_frame": 1842, "start_seconds": 42.44, "end_seconds": 73.72, "duration_seconds": 31.28, "start_time": "00:00:42.440", "end_time": "00:01:13.720", "frame_count": 782, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_001_before_f1060.jpg", "start": "evidence/0014/interval_001_start_f1061.jpg", "end": "evidence/0014/interval_001_end_f1842.jpg", "after": "evidence/0014/interval_001_after_f1843.jpg"}}, {"start_frame": 7887, "end_frame": 8186, "start_seconds": 315.48, "end_seconds": 327.48, "duration_seconds": 12.0, "start_time": "00:05:15.480", "end_time": "00:05:27.480", "frame_count": 300, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_002_before_f7886.jpg", "start": "evidence/0014/interval_002_start_f7887.jpg", "end": "evidence/0014/interval_002_end_f8186.jpg", "after": "evidence/0014/interval_002_after_f8187.jpg"}}, {"start_frame": 9287, "end_frame": 9578, "start_seconds": 371.48, "end_seconds": 383.16, "duration_seconds": 11.68, "start_time": "00:06:11.480", "end_time": "00:06:23.160", "frame_count": 292, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_003_before_f9286.jpg", "start": "evidence/0014/interval_003_start_f9287.jpg", "end": "evidence/0014/interval_003_end_f9578.jpg", "after": "evidence/0014/interval_003_after_f9579.jpg"}}, {"start_frame": 12362, "end_frame": 12576, "start_seconds": 494.48, "end_seconds": 503.08, "duration_seconds": 8.6, "start_time": "00:08:14.480", "end_time": "00:08:23.080", "frame_count": 215, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_004_before_f12361.jpg", "start": "evidence/0014/interval_004_start_f12362.jpg", "end": "evidence/0014/interval_004_end_f12576.jpg", "after": "evidence/0014/interval_004_after_f12577.jpg"}}, {"start_frame": 12932, "end_frame": 13344, "start_seconds": 517.28, "end_seconds": 533.8, "duration_seconds": 16.52, "start_time": "00:08:37.280", "end_time": "00:08:53.800", "frame_count": 413, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_005_before_f12931.jpg", "start": "evidence/0014/interval_005_start_f12932.jpg", "end": "evidence/0014/interval_005_end_f13344.jpg", "after": "evidence/0014/interval_005_after_f13345.jpg"}}, {"start_frame": 14932, "end_frame": 15101, "start_seconds": 597.28, "end_seconds": 604.08, "duration_seconds": 6.8, "start_time": "00:09:57.280", "end_time": "00:10:04.080", "frame_count": 170, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_006_before_f14931.jpg", "start": "evidence/0014/interval_006_start_f14932.jpg", "end": "evidence/0014/interval_006_end_f15101.jpg", "after": "evidence/0014/interval_006_after_f15102.jpg"}}, {"start_frame": 20627, "end_frame": 20781, "start_seconds": 825.08, "end_seconds": 831.28, "duration_seconds": 6.2, "start_time": "00:13:45.080", "end_time": "00:13:51.280", "frame_count": 155, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_007_before_f20626.jpg", "start": "evidence/0014/interval_007_start_f20627.jpg", "end": "evidence/0014/interval_007_end_f20781.jpg", "after": "evidence/0014/interval_007_after_f20782.jpg"}}, {"start_frame": 46852, "end_frame": 47131, "start_seconds": 1874.08, "end_seconds": 1885.28, "duration_seconds": 11.2, "start_time": "00:31:14.080", "end_time": "00:31:25.280", "frame_count": 280, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_008_before_f46851.jpg", "start": "evidence/0014/interval_008_start_f46852.jpg", "end": "evidence/0014/interval_008_end_f47131.jpg", "after": "evidence/0014/interval_008_after_f47132.jpg"}}, {"start_frame": 83957, "end_frame": 84631, "start_seconds": 3358.28, "end_seconds": 3385.28, "duration_seconds": 27.0, "start_time": "00:55:58.280", "end_time": "00:56:25.280", "frame_count": 675, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_009_before_f83956.jpg", "start": "evidence/0014/interval_009_start_f83957.jpg", "end": "evidence/0014/interval_009_end_f84631.jpg", "after": "evidence/0014/interval_009_after_f84632.jpg"}}, {"start_frame": 152532, "end_frame": 152808, "start_seconds": 6101.28, "end_seconds": 6112.36, "duration_seconds": 11.08, "start_time": "01:41:41.280", "end_time": "01:41:52.360", "frame_count": 277, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_010_before_f152531.jpg", "start": "evidence/0014/interval_010_start_f152532.jpg", "end": "evidence/0014/interval_010_end_f152808.jpg", "after": "evidence/0014/interval_010_after_f152809.jpg"}}, {"start_frame": 159772, "end_frame": 159947, "start_seconds": 6390.88, "end_seconds": 6397.92, "duration_seconds": 7.04, "start_time": "01:46:30.880", "end_time": "01:46:37.920", "frame_count": 176, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_011_before_f159771.jpg", "start": "evidence/0014/interval_011_start_f159772.jpg", "end": "evidence/0014/interval_011_end_f159947.jpg", "after": "evidence/0014/interval_011_after_f159948.jpg"}}, {"start_frame": 161457, "end_frame": 161771, "start_seconds": 6458.28, "end_seconds": 6470.88, "duration_seconds": 12.6, "start_time": "01:47:38.280", "end_time": "01:47:50.880", "frame_count": 315, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_012_before_f161456.jpg", "start": "evidence/0014/interval_012_start_f161457.jpg", "end": "evidence/0014/interval_012_end_f161771.jpg", "after": "evidence/0014/interval_012_after_f161772.jpg"}}, {"start_frame": 164407, "end_frame": 164931, "start_seconds": 6576.28, "end_seconds": 6597.28, "duration_seconds": 21.0, "start_time": "01:49:36.280", "end_time": "01:49:57.280", "frame_count": 525, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_013_before_f164406.jpg", "start": "evidence/0014/interval_013_start_f164407.jpg", "end": "evidence/0014/interval_013_end_f164931.jpg", "after": "evidence/0014/interval_013_after_f164932.jpg"}}, {"start_frame": 165100, "end_frame": 165406, "start_seconds": 6604.0, "end_seconds": 6616.28, "duration_seconds": 12.28, "start_time": "01:50:04.000", "end_time": "01:50:16.280", "frame_count": 307, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_014_before_f165099.jpg", "start": "evidence/0014/interval_014_start_f165100.jpg", "end": "evidence/0014/interval_014_end_f165406.jpg", "after": "evidence/0014/interval_014_after_f165407.jpg"}}, {"start_frame": 170881, "end_frame": 171031, "start_seconds": 6835.24, "end_seconds": 6841.28, "duration_seconds": 6.04, "start_time": "01:53:55.240", "end_time": "01:54:01.280", "frame_count": 151, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_015_before_f170880.jpg", "start": "evidence/0014/interval_015_start_f170881.jpg", "end": "evidence/0014/interval_015_end_f171031.jpg", "after": "evidence/0014/interval_015_after_f171032.jpg"}}, {"start_frame": 214275, "end_frame": 256687, "start_seconds": 8571.0, "end_seconds": 10267.52, "duration_seconds": 1696.52, "start_time": "02:22:51.000", "end_time": "02:51:07.520", "frame_count": 42413, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_016_before_f214274.jpg", "start": "evidence/0014/interval_016_start_f214275.jpg", "end": "evidence/0014/interval_016_end_f256687.jpg", "after": "evidence/0014/interval_016_after_f256688.jpg"}}, {"start_frame": 269272, "end_frame": 270063, "start_seconds": 10770.88, "end_seconds": 10802.56, "duration_seconds": 31.68, "start_time": "02:59:30.880", "end_time": "03:00:02.560", "frame_count": 792, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_017_before_f269271.jpg", "start": "evidence/0014/interval_017_start_f269272.jpg", "end": "evidence/0014/interval_017_end_f270063.jpg", "after": "evidence/0014/interval_017_after_f270064.jpg"}}, {"start_frame": 286533, "end_frame": 300706, "start_seconds": 11461.32, "end_seconds": 12028.28, "duration_seconds": 566.96, "start_time": "03:11:01.320", "end_time": "03:20:28.280", "frame_count": 14174, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0014/interval_018_before_f286532.jpg", "start": "evidence/0014/interval_018_start_f286533.jpg", "end": "evidence/0014/interval_018_end_f300706.jpg", "after": "evidence/0014/interval_018_after_f300707.jpg"}}], "fps": 25.0, "frame_count": 304155, "width": 960, "height": 540, "duration_seconds": 12166.2, "decoded_frames": 304155, "coarse_samples": 12167, "coarse_candidate_count": 19, "interval_count": 18, "blue_screen": true, "blue_screen_seconds": 2496.48, "blue_screen_ratio": 0.20519800759481185, "finished_at": "2026-08-25T01:40:55.687713+08:00", "elapsed_seconds": 308.74589816480875, "video_filename": "0014 - Heico - Rektum - 5.avi"} +{"video_id": "0015", "status": "completed", "started_at": "2026-08-25T01:40:55.708814+08:00", "intervals": [{"start_frame": 98, "end_frame": 871, "start_seconds": 3.92, "end_seconds": 34.88, "duration_seconds": 30.96, "start_time": "00:00:03.920", "end_time": "00:00:34.880", "frame_count": 774, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_001_before_f97.jpg", "start": "evidence/0015/interval_001_start_f98.jpg", "end": "evidence/0015/interval_001_end_f871.jpg", "after": "evidence/0015/interval_001_after_f872.jpg"}}, {"start_frame": 1030, "end_frame": 1304, "start_seconds": 41.2, "end_seconds": 52.2, "duration_seconds": 11.0, "start_time": "00:00:41.200", "end_time": "00:00:52.200", "frame_count": 275, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_002_before_f1029.jpg", "start": "evidence/0015/interval_002_start_f1030.jpg", "end": "evidence/0015/interval_002_end_f1304.jpg", "after": "evidence/0015/interval_002_after_f1305.jpg"}}, {"start_frame": 4210, "end_frame": 4569, "start_seconds": 168.4, "end_seconds": 182.8, "duration_seconds": 14.4, "start_time": "00:02:48.400", "end_time": "00:03:02.800", "frame_count": 360, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_003_before_f4209.jpg", "start": "evidence/0015/interval_003_start_f4210.jpg", "end": "evidence/0015/interval_003_end_f4569.jpg", "after": "evidence/0015/interval_003_after_f4570.jpg"}}, {"start_frame": 24926, "end_frame": 25082, "start_seconds": 997.04, "end_seconds": 1003.32, "duration_seconds": 6.28, "start_time": "00:16:37.040", "end_time": "00:16:43.320", "frame_count": 157, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_004_before_f24925.jpg", "start": "evidence/0015/interval_004_start_f24926.jpg", "end": "evidence/0015/interval_004_end_f25082.jpg", "after": "evidence/0015/interval_004_after_f25083.jpg"}}, {"start_frame": 25652, "end_frame": 26954, "start_seconds": 1026.08, "end_seconds": 1078.2, "duration_seconds": 52.12, "start_time": "00:17:06.080", "end_time": "00:17:58.200", "frame_count": 1303, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_005_before_f25651.jpg", "start": "evidence/0015/interval_005_start_f25652.jpg", "end": "evidence/0015/interval_005_end_f26954.jpg", "after": "evidence/0015/interval_005_after_f26955.jpg"}}, {"start_frame": 71242, "end_frame": 71463, "start_seconds": 2849.68, "end_seconds": 2858.56, "duration_seconds": 8.88, "start_time": "00:47:29.680", "end_time": "00:47:38.560", "frame_count": 222, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_006_before_f71241.jpg", "start": "evidence/0015/interval_006_start_f71242.jpg", "end": "evidence/0015/interval_006_end_f71463.jpg", "after": "evidence/0015/interval_006_after_f71464.jpg"}}, {"start_frame": 72854, "end_frame": 73042, "start_seconds": 2914.16, "end_seconds": 2921.72, "duration_seconds": 7.56, "start_time": "00:48:34.160", "end_time": "00:48:41.720", "frame_count": 189, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_007_before_f72853.jpg", "start": "evidence/0015/interval_007_start_f72854.jpg", "end": "evidence/0015/interval_007_end_f73042.jpg", "after": "evidence/0015/interval_007_after_f73043.jpg"}}, {"start_frame": 118161, "end_frame": 118453, "start_seconds": 4726.44, "end_seconds": 4738.16, "duration_seconds": 11.72, "start_time": "01:18:46.440", "end_time": "01:18:58.160", "frame_count": 293, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_008_before_f118160.jpg", "start": "evidence/0015/interval_008_start_f118161.jpg", "end": "evidence/0015/interval_008_end_f118453.jpg", "after": "evidence/0015/interval_008_after_f118454.jpg"}}, {"start_frame": 154118, "end_frame": 154421, "start_seconds": 6164.72, "end_seconds": 6176.88, "duration_seconds": 12.16, "start_time": "01:42:44.720", "end_time": "01:42:56.880", "frame_count": 304, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_009_before_f154117.jpg", "start": "evidence/0015/interval_009_start_f154118.jpg", "end": "evidence/0015/interval_009_end_f154421.jpg", "after": "evidence/0015/interval_009_after_f154422.jpg"}}, {"start_frame": 225395, "end_frame": 225584, "start_seconds": 9015.8, "end_seconds": 9023.4, "duration_seconds": 7.6, "start_time": "02:30:15.800", "end_time": "02:30:23.400", "frame_count": 190, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_010_before_f225394.jpg", "start": "evidence/0015/interval_010_start_f225395.jpg", "end": "evidence/0015/interval_010_end_f225584.jpg", "after": "evidence/0015/interval_010_after_f225585.jpg"}}, {"start_frame": 238226, "end_frame": 238394, "start_seconds": 9529.04, "end_seconds": 9535.8, "duration_seconds": 6.76, "start_time": "02:38:49.040", "end_time": "02:38:55.800", "frame_count": 169, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_011_before_f238225.jpg", "start": "evidence/0015/interval_011_start_f238226.jpg", "end": "evidence/0015/interval_011_end_f238394.jpg", "after": "evidence/0015/interval_011_after_f238395.jpg"}}, {"start_frame": 239171, "end_frame": 239361, "start_seconds": 9566.84, "end_seconds": 9574.48, "duration_seconds": 7.64, "start_time": "02:39:26.840", "end_time": "02:39:34.480", "frame_count": 191, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_012_before_f239170.jpg", "start": "evidence/0015/interval_012_start_f239171.jpg", "end": "evidence/0015/interval_012_end_f239361.jpg", "after": "evidence/0015/interval_012_after_f239362.jpg"}}, {"start_frame": 252669, "end_frame": 252901, "start_seconds": 10106.76, "end_seconds": 10116.08, "duration_seconds": 9.32, "start_time": "02:48:26.760", "end_time": "02:48:36.080", "frame_count": 233, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_013_before_f252668.jpg", "start": "evidence/0015/interval_013_start_f252669.jpg", "end": "evidence/0015/interval_013_end_f252901.jpg", "after": "evidence/0015/interval_013_after_f252902.jpg"}}, {"start_frame": 254945, "end_frame": 369078, "start_seconds": 10197.8, "end_seconds": 14763.16, "duration_seconds": 4565.36, "start_time": "02:49:57.800", "end_time": "04:06:03.160", "frame_count": 114134, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_014_before_f254944.jpg", "start": "evidence/0015/interval_014_start_f254945.jpg", "end": "evidence/0015/interval_014_end_f369078.jpg", "after": "evidence/0015/interval_014_after_f369079.jpg"}}, {"start_frame": 380814, "end_frame": 381016, "start_seconds": 15232.56, "end_seconds": 15240.68, "duration_seconds": 8.12, "start_time": "04:13:52.560", "end_time": "04:14:00.680", "frame_count": 203, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0015/interval_015_before_f380813.jpg", "start": "evidence/0015/interval_015_start_f380814.jpg", "end": "evidence/0015/interval_015_end_f381016.jpg", "after": "evidence/0015/interval_015_after_f381017.jpg"}}], "fps": 25.0, "frame_count": 407549, "width": 960, "height": 540, "duration_seconds": 16301.96, "decoded_frames": 407549, "coarse_samples": 16302, "coarse_candidate_count": 16, "interval_count": 15, "blue_screen": true, "blue_screen_seconds": 4759.879999999999, "blue_screen_ratio": 0.29198206841385943, "finished_at": "2026-08-25T01:48:59.344938+08:00", "elapsed_seconds": 483.63615786843, "video_filename": "0015 - Heico - Rektum - 6.avi"} +{"video_id": "0016", "status": "completed", "started_at": "2026-08-25T01:48:59.364887+08:00", "intervals": [{"start_frame": 286, "end_frame": 504, "start_seconds": 11.44, "end_seconds": 20.2, "duration_seconds": 8.76, "start_time": "00:00:11.440", "end_time": "00:00:20.200", "frame_count": 219, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_001_before_f285.jpg", "start": "evidence/0016/interval_001_start_f286.jpg", "end": "evidence/0016/interval_001_end_f504.jpg", "after": "evidence/0016/interval_001_after_f505.jpg"}}, {"start_frame": 1535, "end_frame": 2679, "start_seconds": 61.4, "end_seconds": 107.2, "duration_seconds": 45.8, "start_time": "00:01:01.400", "end_time": "00:01:47.200", "frame_count": 1145, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_002_before_f1534.jpg", "start": "evidence/0016/interval_002_start_f1535.jpg", "end": "evidence/0016/interval_002_end_f2679.jpg", "after": "evidence/0016/interval_002_after_f2680.jpg"}}, {"start_frame": 3366, "end_frame": 3528, "start_seconds": 134.64, "end_seconds": 141.16, "duration_seconds": 6.52, "start_time": "00:02:14.640", "end_time": "00:02:21.160", "frame_count": 163, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_003_before_f3365.jpg", "start": "evidence/0016/interval_003_start_f3366.jpg", "end": "evidence/0016/interval_003_end_f3528.jpg", "after": "evidence/0016/interval_003_after_f3529.jpg"}}, {"start_frame": 76469, "end_frame": 77327, "start_seconds": 3058.76, "end_seconds": 3093.12, "duration_seconds": 34.36, "start_time": "00:50:58.760", "end_time": "00:51:33.120", "frame_count": 859, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_004_before_f76468.jpg", "start": "evidence/0016/interval_004_start_f76469.jpg", "end": "evidence/0016/interval_004_end_f77327.jpg", "after": "evidence/0016/interval_004_after_f77328.jpg"}}, {"start_frame": 111617, "end_frame": 111991, "start_seconds": 4464.68, "end_seconds": 4479.68, "duration_seconds": 15.0, "start_time": "01:14:24.680", "end_time": "01:14:39.680", "frame_count": 375, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_005_before_f111616.jpg", "start": "evidence/0016/interval_005_start_f111617.jpg", "end": "evidence/0016/interval_005_end_f111991.jpg", "after": "evidence/0016/interval_005_after_f111992.jpg"}}, {"start_frame": 176683, "end_frame": 226219, "start_seconds": 7067.32, "end_seconds": 9048.8, "duration_seconds": 1981.48, "start_time": "01:57:47.320", "end_time": "02:30:48.800", "frame_count": 49537, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_006_before_f176682.jpg", "start": "evidence/0016/interval_006_start_f176683.jpg", "end": "evidence/0016/interval_006_end_f226219.jpg", "after": "evidence/0016/interval_006_after_f226220.jpg"}}, {"start_frame": 230857, "end_frame": 231447, "start_seconds": 9234.28, "end_seconds": 9257.92, "duration_seconds": 23.64, "start_time": "02:33:54.280", "end_time": "02:34:17.920", "frame_count": 591, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_007_before_f230856.jpg", "start": "evidence/0016/interval_007_start_f230857.jpg", "end": "evidence/0016/interval_007_end_f231447.jpg", "after": "evidence/0016/interval_007_after_f231448.jpg"}}, {"start_frame": 245796, "end_frame": 246039, "start_seconds": 9831.84, "end_seconds": 9841.6, "duration_seconds": 9.76, "start_time": "02:43:51.840", "end_time": "02:44:01.600", "frame_count": 244, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_008_before_f245795.jpg", "start": "evidence/0016/interval_008_start_f245796.jpg", "end": "evidence/0016/interval_008_end_f246039.jpg", "after": "evidence/0016/interval_008_after_f246040.jpg"}}, {"start_frame": 256669, "end_frame": 257081, "start_seconds": 10266.76, "end_seconds": 10283.28, "duration_seconds": 16.52, "start_time": "02:51:06.760", "end_time": "02:51:23.280", "frame_count": 413, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_009_before_f256668.jpg", "start": "evidence/0016/interval_009_start_f256669.jpg", "end": "evidence/0016/interval_009_end_f257081.jpg", "after": "evidence/0016/interval_009_after_f257082.jpg"}}, {"start_frame": 262393, "end_frame": 262789, "start_seconds": 10495.72, "end_seconds": 10511.6, "duration_seconds": 15.88, "start_time": "02:54:55.720", "end_time": "02:55:11.600", "frame_count": 397, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_010_before_f262392.jpg", "start": "evidence/0016/interval_010_start_f262393.jpg", "end": "evidence/0016/interval_010_end_f262789.jpg", "after": "evidence/0016/interval_010_after_f262790.jpg"}}, {"start_frame": 273830, "end_frame": 274150, "start_seconds": 10953.2, "end_seconds": 10966.04, "duration_seconds": 12.84, "start_time": "03:02:33.200", "end_time": "03:02:46.040", "frame_count": 321, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_011_before_f273829.jpg", "start": "evidence/0016/interval_011_start_f273830.jpg", "end": "evidence/0016/interval_011_end_f274150.jpg", "after": "evidence/0016/interval_011_after_f274151.jpg"}}, {"start_frame": 275856, "end_frame": 276843, "start_seconds": 11034.24, "end_seconds": 11073.76, "duration_seconds": 39.52, "start_time": "03:03:54.240", "end_time": "03:04:33.760", "frame_count": 988, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_012_before_f275855.jpg", "start": "evidence/0016/interval_012_start_f275856.jpg", "end": "evidence/0016/interval_012_end_f276843.jpg", "after": "evidence/0016/interval_012_after_f276844.jpg"}}, {"start_frame": 306912, "end_frame": 317030, "start_seconds": 12276.48, "end_seconds": 12681.24, "duration_seconds": 404.76, "start_time": "03:24:36.480", "end_time": "03:31:21.240", "frame_count": 10119, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_013_before_f306911.jpg", "start": "evidence/0016/interval_013_start_f306912.jpg", "end": "evidence/0016/interval_013_end_f317030.jpg", "after": "evidence/0016/interval_013_after_f317031.jpg"}}, {"start_frame": 325605, "end_frame": 337482, "start_seconds": 13024.2, "end_seconds": 13499.32, "duration_seconds": 475.12, "start_time": "03:37:04.200", "end_time": "03:44:59.320", "frame_count": 11878, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_014_before_f325604.jpg", "start": "evidence/0016/interval_014_start_f325605.jpg", "end": "evidence/0016/interval_014_end_f337482.jpg", "after": "evidence/0016/interval_014_after_f337483.jpg"}}, {"start_frame": 337785, "end_frame": 349126, "start_seconds": 13511.4, "end_seconds": 13965.08, "duration_seconds": 453.68, "start_time": "03:45:11.400", "end_time": "03:52:45.080", "frame_count": 11342, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_015_before_f337784.jpg", "start": "evidence/0016/interval_015_start_f337785.jpg", "end": "evidence/0016/interval_015_end_f349126.jpg", "after": "evidence/0016/interval_015_after_f349127.jpg"}}, {"start_frame": 380708, "end_frame": 381065, "start_seconds": 15228.32, "end_seconds": 15242.64, "duration_seconds": 14.32, "start_time": "04:13:48.320", "end_time": "04:14:02.640", "frame_count": 358, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0016/interval_016_before_f380707.jpg", "start": "evidence/0016/interval_016_start_f380708.jpg", "end": "evidence/0016/interval_016_end_f381065.jpg", "after": "evidence/0016/interval_016_after_f381066.jpg"}}], "fps": 25.0, "frame_count": 391223, "width": 960, "height": 540, "duration_seconds": 15648.92, "decoded_frames": 391223, "coarse_samples": 15649, "coarse_candidate_count": 17, "interval_count": 16, "blue_screen": true, "blue_screen_seconds": 3557.96, "blue_screen_ratio": 0.22736137701515505, "finished_at": "2026-08-25T01:58:11.917684+08:00", "elapsed_seconds": 552.5528265405446, "video_filename": "0016 - Heico - Rektum - 7.avi"} +{"video_id": "0017", "status": "completed", "started_at": "2026-08-25T01:58:12.005132+08:00", "intervals": [{"start_frame": 3448, "end_frame": 3630, "start_seconds": 137.92, "end_seconds": 145.24, "duration_seconds": 7.32, "start_time": "00:02:17.920", "end_time": "00:02:25.240", "frame_count": 183, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0017/interval_001_before_f3447.jpg", "start": "evidence/0017/interval_001_start_f3448.jpg", "end": "evidence/0017/interval_001_end_f3630.jpg", "after": "evidence/0017/interval_001_after_f3631.jpg"}}, {"start_frame": 3746, "end_frame": 4105, "start_seconds": 149.84, "end_seconds": 164.24, "duration_seconds": 14.4, "start_time": "00:02:29.840", "end_time": "00:02:44.240", "frame_count": 360, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0017/interval_002_before_f3745.jpg", "start": "evidence/0017/interval_002_start_f3746.jpg", "end": "evidence/0017/interval_002_end_f4105.jpg", "after": "evidence/0017/interval_002_after_f4106.jpg"}}, {"start_frame": 5012, "end_frame": 5189, "start_seconds": 200.48, "end_seconds": 207.6, "duration_seconds": 7.12, "start_time": "00:03:20.480", "end_time": "00:03:27.600", "frame_count": 178, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0017/interval_003_before_f5011.jpg", "start": "evidence/0017/interval_003_start_f5012.jpg", "end": "evidence/0017/interval_003_end_f5189.jpg", "after": "evidence/0017/interval_003_after_f5190.jpg"}}, {"start_frame": 10783, "end_frame": 11192, "start_seconds": 431.32, "end_seconds": 447.72, "duration_seconds": 16.4, "start_time": "00:07:11.320", "end_time": "00:07:27.720", "frame_count": 410, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0017/interval_004_before_f10782.jpg", "start": "evidence/0017/interval_004_start_f10783.jpg", "end": "evidence/0017/interval_004_end_f11192.jpg", "after": "evidence/0017/interval_004_after_f11193.jpg"}}, {"start_frame": 143141, "end_frame": 185332, "start_seconds": 5725.64, "end_seconds": 7413.32, "duration_seconds": 1687.68, "start_time": "01:35:25.640", "end_time": "02:03:33.320", "frame_count": 42192, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0017/interval_005_before_f143140.jpg", "start": "evidence/0017/interval_005_start_f143141.jpg", "end": "evidence/0017/interval_005_end_f185332.jpg", "after": "evidence/0017/interval_005_after_f185333.jpg"}}, {"start_frame": 210088, "end_frame": 210262, "start_seconds": 8403.52, "end_seconds": 8410.52, "duration_seconds": 7.0, "start_time": "02:20:03.520", "end_time": "02:20:10.520", "frame_count": 175, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0017/interval_006_before_f210087.jpg", "start": "evidence/0017/interval_006_start_f210088.jpg", "end": "evidence/0017/interval_006_end_f210262.jpg", "after": "evidence/0017/interval_006_after_f210263.jpg"}}], "fps": 25.0, "frame_count": 257801, "width": 960, "height": 540, "duration_seconds": 10312.04, "decoded_frames": 257801, "coarse_samples": 10313, "coarse_candidate_count": 6, "interval_count": 6, "blue_screen": true, "blue_screen_seconds": 1739.92, "blue_screen_ratio": 0.1687270414001497, "finished_at": "2026-08-25T02:01:54.842206+08:00", "elapsed_seconds": 222.83711133897305, "video_filename": "0017 - Heico - Rektum - 8.avi"} +{"video_id": "0018", "status": "completed", "started_at": "2026-08-25T02:01:54.859491+08:00", "intervals": [{"start_frame": 4605, "end_frame": 6290, "start_seconds": 184.2, "end_seconds": 251.64, "duration_seconds": 67.44, "start_time": "00:03:04.200", "end_time": "00:04:11.640", "frame_count": 1686, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0018/interval_001_before_f4604.jpg", "start": "evidence/0018/interval_001_start_f4605.jpg", "end": "evidence/0018/interval_001_end_f6290.jpg", "after": "evidence/0018/interval_001_after_f6291.jpg"}}, {"start_frame": 196494, "end_frame": 196665, "start_seconds": 7859.76, "end_seconds": 7866.64, "duration_seconds": 6.88, "start_time": "02:10:59.760", "end_time": "02:11:06.640", "frame_count": 172, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0018/interval_002_before_f196493.jpg", "start": "evidence/0018/interval_002_start_f196494.jpg", "end": "evidence/0018/interval_002_end_f196665.jpg", "after": "evidence/0018/interval_002_after_f196666.jpg"}}, {"start_frame": 207338, "end_frame": 260298, "start_seconds": 8293.52, "end_seconds": 10411.96, "duration_seconds": 2118.44, "start_time": "02:18:13.520", "end_time": "02:53:31.960", "frame_count": 52961, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0018/interval_003_before_f207337.jpg", "start": "evidence/0018/interval_003_start_f207338.jpg", "end": "evidence/0018/interval_003_end_f260298.jpg", "after": "evidence/0018/interval_003_after_f260299.jpg"}}, {"start_frame": 267668, "end_frame": 268079, "start_seconds": 10706.72, "end_seconds": 10723.2, "duration_seconds": 16.48, "start_time": "02:58:26.720", "end_time": "02:58:43.200", "frame_count": 412, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0018/interval_004_before_f267667.jpg", "start": "evidence/0018/interval_004_start_f267668.jpg", "end": "evidence/0018/interval_004_end_f268079.jpg", "after": "evidence/0018/interval_004_after_f268080.jpg"}}, {"start_frame": 269857, "end_frame": 270216, "start_seconds": 10794.28, "end_seconds": 10808.68, "duration_seconds": 14.4, "start_time": "02:59:54.280", "end_time": "03:00:08.680", "frame_count": 360, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0018/interval_005_before_f269856.jpg", "start": "evidence/0018/interval_005_start_f269857.jpg", "end": "evidence/0018/interval_005_end_f270216.jpg", "after": "evidence/0018/interval_005_after_f270217.jpg"}}, {"start_frame": 345674, "end_frame": 346337, "start_seconds": 13826.96, "end_seconds": 13853.52, "duration_seconds": 26.56, "start_time": "03:50:26.960", "end_time": "03:50:53.520", "frame_count": 664, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0018/interval_006_before_f345673.jpg", "start": "evidence/0018/interval_006_start_f345674.jpg", "end": "evidence/0018/interval_006_end_f346337.jpg", "after": "evidence/0018/interval_006_after_f346338.jpg"}}], "fps": 25.0, "frame_count": 356815, "width": 960, "height": 540, "duration_seconds": 14272.6, "decoded_frames": 356815, "coarse_samples": 14273, "coarse_candidate_count": 6, "interval_count": 6, "blue_screen": true, "blue_screen_seconds": 2250.2000000000003, "blue_screen_ratio": 0.1576587307147962, "finished_at": "2026-08-25T02:07:01.589210+08:00", "elapsed_seconds": 306.7297497782856, "video_filename": "0018 - Heico - Rektum - 9.avi"} +{"video_id": "0019", "status": "completed", "started_at": "2026-08-25T02:07:01.609703+08:00", "intervals": [{"start_frame": 88458, "end_frame": 88727, "start_seconds": 3538.32, "end_seconds": 3549.12, "duration_seconds": 10.8, "start_time": "00:58:58.320", "end_time": "00:59:09.120", "frame_count": 270, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0019/interval_001_before_f88457.jpg", "start": "evidence/0019/interval_001_start_f88458.jpg", "end": "evidence/0019/interval_001_end_f88727.jpg", "after": "evidence/0019/interval_001_after_f88728.jpg"}}, {"start_frame": 101634, "end_frame": 155868, "start_seconds": 4065.36, "end_seconds": 6234.76, "duration_seconds": 2169.4, "start_time": "01:07:45.360", "end_time": "01:43:54.760", "frame_count": 54235, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0019/interval_002_before_f101633.jpg", "start": "evidence/0019/interval_002_start_f101634.jpg", "end": "evidence/0019/interval_002_end_f155868.jpg", "after": "evidence/0019/interval_002_after_f155869.jpg"}}, {"start_frame": 157074, "end_frame": 157625, "start_seconds": 6282.96, "end_seconds": 6305.04, "duration_seconds": 22.08, "start_time": "01:44:42.960", "end_time": "01:45:05.040", "frame_count": 552, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0019/interval_003_before_f157073.jpg", "start": "evidence/0019/interval_003_start_f157074.jpg", "end": "evidence/0019/interval_003_end_f157625.jpg", "after": "evidence/0019/interval_003_after_f157626.jpg"}}], "fps": 25.0, "frame_count": 179116, "width": 960, "height": 540, "duration_seconds": 7164.64, "decoded_frames": 179116, "coarse_samples": 7165, "coarse_candidate_count": 3, "interval_count": 3, "blue_screen": true, "blue_screen_seconds": 2202.28, "blue_screen_ratio": 0.3073818084369906, "finished_at": "2026-08-25T02:10:26.430892+08:00", "elapsed_seconds": 204.82121763564646, "video_filename": "0019 - Heico - Rektum - 10.avi"} +{"video_id": "0020", "status": "completed", "started_at": "2026-08-25T02:10:26.447833+08:00", "intervals": [{"start_frame": 76741, "end_frame": 120650, "start_seconds": 3069.64, "end_seconds": 4826.04, "duration_seconds": 1756.4, "start_time": "00:51:09.640", "end_time": "01:20:26.040", "frame_count": 43910, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0020/interval_001_before_f76740.jpg", "start": "evidence/0020/interval_001_start_f76741.jpg", "end": "evidence/0020/interval_001_end_f120650.jpg", "after": "evidence/0020/interval_001_after_f120651.jpg"}}], "fps": 25.0, "frame_count": 145377, "width": 960, "height": 540, "duration_seconds": 5815.08, "decoded_frames": 145377, "coarse_samples": 5816, "coarse_candidate_count": 1, "interval_count": 1, "blue_screen": true, "blue_screen_seconds": 1756.4, "blue_screen_ratio": 0.30204227628854635, "finished_at": "2026-08-25T02:13:05.994141+08:00", "elapsed_seconds": 159.5463363789022, "video_filename": "0020 - Heico - Sigma - 1.avi"} +{"video_id": "0021", "status": "completed", "started_at": "2026-08-25T02:13:06.012766+08:00", "intervals": [{"start_frame": 1403, "end_frame": 1602, "start_seconds": 56.12, "end_seconds": 64.12, "duration_seconds": 8.0, "start_time": "00:00:56.120", "end_time": "00:01:04.120", "frame_count": 200, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0021/interval_001_before_f1402.jpg", "start": "evidence/0021/interval_001_start_f1403.jpg", "end": "evidence/0021/interval_001_end_f1602.jpg", "after": "evidence/0021/interval_001_after_f1603.jpg"}}, {"start_frame": 75202, "end_frame": 129434, "start_seconds": 3008.08, "end_seconds": 5177.4, "duration_seconds": 2169.32, "start_time": "00:50:08.080", "end_time": "01:26:17.400", "frame_count": 54233, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0021/interval_002_before_f75201.jpg", "start": "evidence/0021/interval_002_start_f75202.jpg", "end": "evidence/0021/interval_002_end_f129434.jpg", "after": "evidence/0021/interval_002_after_f129435.jpg"}}, {"start_frame": 142160, "end_frame": 142314, "start_seconds": 5686.4, "end_seconds": 5692.6, "duration_seconds": 6.2, "start_time": "01:34:46.400", "end_time": "01:34:52.600", "frame_count": 155, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0021/interval_003_before_f142159.jpg", "start": "evidence/0021/interval_003_start_f142160.jpg", "end": "evidence/0021/interval_003_end_f142314.jpg", "after": "evidence/0021/interval_003_after_f142315.jpg"}}, {"start_frame": 145460, "end_frame": 145669, "start_seconds": 5818.4, "end_seconds": 5826.8, "duration_seconds": 8.4, "start_time": "01:36:58.400", "end_time": "01:37:06.800", "frame_count": 210, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0021/interval_004_before_f145459.jpg", "start": "evidence/0021/interval_004_start_f145460.jpg", "end": "evidence/0021/interval_004_end_f145669.jpg", "after": "evidence/0021/interval_004_after_f145670.jpg"}}], "fps": 25.0, "frame_count": 155037, "width": 960, "height": 540, "duration_seconds": 6201.48, "decoded_frames": 155037, "coarse_samples": 6202, "coarse_candidate_count": 4, "interval_count": 4, "blue_screen": true, "blue_screen_seconds": 2191.92, "blue_screen_ratio": 0.3534511116701175, "finished_at": "2026-08-25T02:15:59.983928+08:00", "elapsed_seconds": 173.971192330122, "video_filename": "0021 - Heico - Sigma - 2.avi"} +{"video_id": "0022", "status": "completed", "started_at": "2026-08-25T02:16:00.041849+08:00", "intervals": [{"start_frame": 4196, "end_frame": 4772, "start_seconds": 167.84, "end_seconds": 190.92, "duration_seconds": 23.08, "start_time": "00:02:47.840", "end_time": "00:03:10.920", "frame_count": 577, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_001_before_f4195.jpg", "start": "evidence/0022/interval_001_start_f4196.jpg", "end": "evidence/0022/interval_001_end_f4772.jpg", "after": "evidence/0022/interval_001_after_f4773.jpg"}}, {"start_frame": 58606, "end_frame": 59151, "start_seconds": 2344.24, "end_seconds": 2366.08, "duration_seconds": 21.84, "start_time": "00:39:04.240", "end_time": "00:39:26.080", "frame_count": 546, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_002_before_f58605.jpg", "start": "evidence/0022/interval_002_start_f58606.jpg", "end": "evidence/0022/interval_002_end_f59151.jpg", "after": "evidence/0022/interval_002_after_f59152.jpg"}}, {"start_frame": 108376, "end_frame": 109561, "start_seconds": 4335.04, "end_seconds": 4382.48, "duration_seconds": 47.44, "start_time": "01:12:15.040", "end_time": "01:13:02.480", "frame_count": 1186, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_003_before_f108375.jpg", "start": "evidence/0022/interval_003_start_f108376.jpg", "end": "evidence/0022/interval_003_end_f109561.jpg", "after": "evidence/0022/interval_003_after_f109562.jpg"}}, {"start_frame": 150566, "end_frame": 151065, "start_seconds": 6022.64, "end_seconds": 6042.64, "duration_seconds": 20.0, "start_time": "01:40:22.640", "end_time": "01:40:42.640", "frame_count": 500, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_004_before_f150565.jpg", "start": "evidence/0022/interval_004_start_f150566.jpg", "end": "evidence/0022/interval_004_end_f151065.jpg", "after": "evidence/0022/interval_004_after_f151066.jpg"}}, {"start_frame": 213564, "end_frame": 213819, "start_seconds": 8542.56, "end_seconds": 8552.8, "duration_seconds": 10.24, "start_time": "02:22:22.560", "end_time": "02:22:32.800", "frame_count": 256, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_005_before_f213563.jpg", "start": "evidence/0022/interval_005_start_f213564.jpg", "end": "evidence/0022/interval_005_end_f213819.jpg", "after": "evidence/0022/interval_005_after_f213820.jpg"}}, {"start_frame": 213927, "end_frame": 214506, "start_seconds": 8557.08, "end_seconds": 8580.28, "duration_seconds": 23.2, "start_time": "02:22:37.080", "end_time": "02:23:00.280", "frame_count": 580, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_006_before_f213926.jpg", "start": "evidence/0022/interval_006_start_f213927.jpg", "end": "evidence/0022/interval_006_end_f214506.jpg", "after": "evidence/0022/interval_006_after_f214507.jpg"}}, {"start_frame": 235981, "end_frame": 292296, "start_seconds": 9439.24, "end_seconds": 11691.88, "duration_seconds": 2252.64, "start_time": "02:37:19.240", "end_time": "03:14:51.880", "frame_count": 56316, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_007_before_f235980.jpg", "start": "evidence/0022/interval_007_start_f235981.jpg", "end": "evidence/0022/interval_007_end_f292296.jpg", "after": "evidence/0022/interval_007_after_f292297.jpg"}}, {"start_frame": 308248, "end_frame": 308950, "start_seconds": 12329.92, "end_seconds": 12358.04, "duration_seconds": 28.12, "start_time": "03:25:29.920", "end_time": "03:25:58.040", "frame_count": 703, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_008_before_f308247.jpg", "start": "evidence/0022/interval_008_start_f308248.jpg", "end": "evidence/0022/interval_008_end_f308950.jpg", "after": "evidence/0022/interval_008_after_f308951.jpg"}}, {"start_frame": 313811, "end_frame": 314209, "start_seconds": 12552.44, "end_seconds": 12568.4, "duration_seconds": 15.96, "start_time": "03:29:12.440", "end_time": "03:29:28.400", "frame_count": 399, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_009_before_f313810.jpg", "start": "evidence/0022/interval_009_start_f313811.jpg", "end": "evidence/0022/interval_009_end_f314209.jpg", "after": "evidence/0022/interval_009_after_f314210.jpg"}}, {"start_frame": 331328, "end_frame": 332053, "start_seconds": 13253.12, "end_seconds": 13282.16, "duration_seconds": 29.04, "start_time": "03:40:53.120", "end_time": "03:41:22.160", "frame_count": 726, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0022/interval_010_before_f331327.jpg", "start": "evidence/0022/interval_010_start_f331328.jpg", "end": "evidence/0022/interval_010_end_f332053.jpg"}}], "fps": 25.0, "frame_count": 332054, "width": 960, "height": 540, "duration_seconds": 13282.16, "decoded_frames": 332054, "coarse_samples": 13283, "coarse_candidate_count": 10, "interval_count": 10, "blue_screen": true, "blue_screen_seconds": 2471.56, "blue_screen_ratio": 0.18608117956717882, "finished_at": "2026-08-25T02:21:48.827399+08:00", "elapsed_seconds": 348.7855956647545, "video_filename": "0022 - Heico - Sigma - 3.avi"} +{"video_id": "0023", "status": "completed", "started_at": "2026-08-25T02:21:48.852812+08:00", "intervals": [{"start_frame": 22393, "end_frame": 22586, "start_seconds": 895.72, "end_seconds": 903.48, "duration_seconds": 7.76, "start_time": "00:14:55.720", "end_time": "00:15:03.480", "frame_count": 194, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0023/interval_001_before_f22392.jpg", "start": "evidence/0023/interval_001_start_f22393.jpg", "end": "evidence/0023/interval_001_end_f22586.jpg", "after": "evidence/0023/interval_001_after_f22587.jpg"}}, {"start_frame": 23387, "end_frame": 23797, "start_seconds": 935.48, "end_seconds": 951.92, "duration_seconds": 16.44, "start_time": "00:15:35.480", "end_time": "00:15:51.920", "frame_count": 411, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0023/interval_002_before_f23386.jpg", "start": "evidence/0023/interval_002_start_f23387.jpg", "end": "evidence/0023/interval_002_end_f23797.jpg", "after": "evidence/0023/interval_002_after_f23798.jpg"}}, {"start_frame": 58098, "end_frame": 58324, "start_seconds": 2323.92, "end_seconds": 2333.0, "duration_seconds": 9.08, "start_time": "00:38:43.920", "end_time": "00:38:53.000", "frame_count": 227, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0023/interval_003_before_f58097.jpg", "start": "evidence/0023/interval_003_start_f58098.jpg", "end": "evidence/0023/interval_003_end_f58324.jpg", "after": "evidence/0023/interval_003_after_f58325.jpg"}}, {"start_frame": 89007, "end_frame": 89196, "start_seconds": 3560.28, "end_seconds": 3567.88, "duration_seconds": 7.6, "start_time": "00:59:20.280", "end_time": "00:59:27.880", "frame_count": 190, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0023/interval_004_before_f89006.jpg", "start": "evidence/0023/interval_004_start_f89007.jpg", "end": "evidence/0023/interval_004_end_f89196.jpg", "after": "evidence/0023/interval_004_after_f89197.jpg"}}, {"start_frame": 135563, "end_frame": 136671, "start_seconds": 5422.52, "end_seconds": 5466.88, "duration_seconds": 44.36, "start_time": "01:30:22.520", "end_time": "01:31:06.880", "frame_count": 1109, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0023/interval_005_before_f135562.jpg", "start": "evidence/0023/interval_005_start_f135563.jpg", "end": "evidence/0023/interval_005_end_f136671.jpg", "after": "evidence/0023/interval_005_after_f136672.jpg"}}, {"start_frame": 136822, "end_frame": 188918, "start_seconds": 5472.88, "end_seconds": 7556.76, "duration_seconds": 2083.88, "start_time": "01:31:12.880", "end_time": "02:05:56.760", "frame_count": 52097, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0023/interval_006_before_f136821.jpg", "start": "evidence/0023/interval_006_start_f136822.jpg", "end": "evidence/0023/interval_006_end_f188918.jpg", "after": "evidence/0023/interval_006_after_f188919.jpg"}}, {"start_frame": 189491, "end_frame": 189698, "start_seconds": 7579.64, "end_seconds": 7587.96, "duration_seconds": 8.32, "start_time": "02:06:19.640", "end_time": "02:06:27.960", "frame_count": 208, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0023/interval_007_before_f189490.jpg", "start": "evidence/0023/interval_007_start_f189491.jpg", "end": "evidence/0023/interval_007_end_f189698.jpg", "after": "evidence/0023/interval_007_after_f189699.jpg"}}, {"start_frame": 216348, "end_frame": 216559, "start_seconds": 8653.92, "end_seconds": 8662.4, "duration_seconds": 8.48, "start_time": "02:24:13.920", "end_time": "02:24:22.400", "frame_count": 212, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0023/interval_008_before_f216347.jpg", "start": "evidence/0023/interval_008_start_f216348.jpg", "end": "evidence/0023/interval_008_end_f216559.jpg", "after": "evidence/0023/interval_008_after_f216560.jpg"}}], "fps": 25.0, "frame_count": 220054, "width": 960, "height": 540, "duration_seconds": 8802.16, "decoded_frames": 220054, "coarse_samples": 8803, "coarse_candidate_count": 9, "interval_count": 8, "blue_screen": true, "blue_screen_seconds": 2185.92, "blue_screen_ratio": 0.24833904405282342, "finished_at": "2026-08-25T02:25:25.549753+08:00", "elapsed_seconds": 216.69697524048388, "video_filename": "0023 - Heico - Sigma - 4.avi"} +{"video_id": "0024", "status": "completed", "started_at": "2026-08-25T02:25:25.572049+08:00", "intervals": [{"start_frame": 3523, "end_frame": 3791, "start_seconds": 140.92, "end_seconds": 151.68, "duration_seconds": 10.76, "start_time": "00:02:20.920", "end_time": "00:02:31.680", "frame_count": 269, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0024/interval_001_before_f3522.jpg", "start": "evidence/0024/interval_001_start_f3523.jpg", "end": "evidence/0024/interval_001_end_f3791.jpg", "after": "evidence/0024/interval_001_after_f3792.jpg"}}, {"start_frame": 49240, "end_frame": 49671, "start_seconds": 1969.6, "end_seconds": 1986.88, "duration_seconds": 17.28, "start_time": "00:32:49.600", "end_time": "00:33:06.880", "frame_count": 432, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0024/interval_002_before_f49239.jpg", "start": "evidence/0024/interval_002_start_f49240.jpg", "end": "evidence/0024/interval_002_end_f49671.jpg", "after": "evidence/0024/interval_002_after_f49672.jpg"}}, {"start_frame": 72121, "end_frame": 129067, "start_seconds": 2884.84, "end_seconds": 5162.72, "duration_seconds": 2277.88, "start_time": "00:48:04.840", "end_time": "01:26:02.720", "frame_count": 56947, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0024/interval_003_before_f72120.jpg", "start": "evidence/0024/interval_003_start_f72121.jpg", "end": "evidence/0024/interval_003_end_f129067.jpg", "after": "evidence/0024/interval_003_after_f129068.jpg"}}], "fps": 25.0, "frame_count": 169644, "width": 960, "height": 540, "duration_seconds": 6785.76, "decoded_frames": 169644, "coarse_samples": 6786, "coarse_candidate_count": 4, "interval_count": 3, "blue_screen": true, "blue_screen_seconds": 2305.92, "blue_screen_ratio": 0.3398175001768409, "finished_at": "2026-08-25T02:28:44.250383+08:00", "elapsed_seconds": 198.6783648133278, "video_filename": "0024 - Heico - Sigma - 5.avi"} +{"video_id": "0025", "status": "completed", "started_at": "2026-08-25T02:28:44.267663+08:00", "intervals": [{"start_frame": 769, "end_frame": 971, "start_seconds": 30.76, "end_seconds": 38.88, "duration_seconds": 8.12, "start_time": "00:00:30.760", "end_time": "00:00:38.880", "frame_count": 203, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_001_before_f768.jpg", "start": "evidence/0025/interval_001_start_f769.jpg", "end": "evidence/0025/interval_001_end_f971.jpg", "after": "evidence/0025/interval_001_after_f972.jpg"}}, {"start_frame": 5809, "end_frame": 5986, "start_seconds": 232.36, "end_seconds": 239.48, "duration_seconds": 7.12, "start_time": "00:03:52.360", "end_time": "00:03:59.480", "frame_count": 178, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_002_before_f5808.jpg", "start": "evidence/0025/interval_002_start_f5809.jpg", "end": "evidence/0025/interval_002_end_f5986.jpg", "after": "evidence/0025/interval_002_after_f5987.jpg"}}, {"start_frame": 53622, "end_frame": 54258, "start_seconds": 2144.88, "end_seconds": 2170.36, "duration_seconds": 25.48, "start_time": "00:35:44.880", "end_time": "00:36:10.360", "frame_count": 637, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_003_before_f53621.jpg", "start": "evidence/0025/interval_003_start_f53622.jpg", "end": "evidence/0025/interval_003_end_f54258.jpg", "after": "evidence/0025/interval_003_after_f54259.jpg"}}, {"start_frame": 96792, "end_frame": 96976, "start_seconds": 3871.68, "end_seconds": 3879.08, "duration_seconds": 7.4, "start_time": "01:04:31.680", "end_time": "01:04:39.080", "frame_count": 185, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_004_before_f96791.jpg", "start": "evidence/0025/interval_004_start_f96792.jpg", "end": "evidence/0025/interval_004_end_f96976.jpg", "after": "evidence/0025/interval_004_after_f96977.jpg"}}, {"start_frame": 119669, "end_frame": 218770, "start_seconds": 4786.76, "end_seconds": 8750.84, "duration_seconds": 3964.08, "start_time": "01:19:46.760", "end_time": "02:25:50.840", "frame_count": 99102, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_005_before_f119668.jpg", "start": "evidence/0025/interval_005_start_f119669.jpg", "end": "evidence/0025/interval_005_end_f218770.jpg", "after": "evidence/0025/interval_005_after_f218771.jpg"}}, {"start_frame": 228355, "end_frame": 230160, "start_seconds": 9134.2, "end_seconds": 9206.44, "duration_seconds": 72.24, "start_time": "02:32:14.200", "end_time": "02:33:26.440", "frame_count": 1806, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_006_before_f228354.jpg", "start": "evidence/0025/interval_006_start_f228355.jpg", "end": "evidence/0025/interval_006_end_f230160.jpg", "after": "evidence/0025/interval_006_after_f230161.jpg"}}, {"start_frame": 250996, "end_frame": 252051, "start_seconds": 10039.84, "end_seconds": 10082.08, "duration_seconds": 42.24, "start_time": "02:47:19.840", "end_time": "02:48:02.080", "frame_count": 1056, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_007_before_f250995.jpg", "start": "evidence/0025/interval_007_start_f250996.jpg", "end": "evidence/0025/interval_007_end_f252051.jpg", "after": "evidence/0025/interval_007_after_f252052.jpg"}}, {"start_frame": 254176, "end_frame": 254633, "start_seconds": 10167.04, "end_seconds": 10185.36, "duration_seconds": 18.32, "start_time": "02:49:27.040", "end_time": "02:49:45.360", "frame_count": 458, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_008_before_f254175.jpg", "start": "evidence/0025/interval_008_start_f254176.jpg", "end": "evidence/0025/interval_008_end_f254633.jpg", "after": "evidence/0025/interval_008_after_f254634.jpg"}}, {"start_frame": 294134, "end_frame": 294341, "start_seconds": 11765.36, "end_seconds": 11773.68, "duration_seconds": 8.32, "start_time": "03:16:05.360", "end_time": "03:16:13.680", "frame_count": 208, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_009_before_f294133.jpg", "start": "evidence/0025/interval_009_start_f294134.jpg", "end": "evidence/0025/interval_009_end_f294341.jpg", "after": "evidence/0025/interval_009_after_f294342.jpg"}}, {"start_frame": 305932, "end_frame": 306158, "start_seconds": 12237.28, "end_seconds": 12246.36, "duration_seconds": 9.08, "start_time": "03:23:57.280", "end_time": "03:24:06.360", "frame_count": 227, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0025/interval_010_before_f305931.jpg", "start": "evidence/0025/interval_010_start_f305932.jpg", "end": "evidence/0025/interval_010_end_f306158.jpg", "after": "evidence/0025/interval_010_after_f306159.jpg"}}], "fps": 25.0, "frame_count": 318286, "width": 960, "height": 540, "duration_seconds": 12731.44, "decoded_frames": 318286, "coarse_samples": 12732, "coarse_candidate_count": 10, "interval_count": 10, "blue_screen": true, "blue_screen_seconds": 4162.399999999999, "blue_screen_ratio": 0.3269386652256146, "finished_at": "2026-08-25T02:34:45.276959+08:00", "elapsed_seconds": 361.0093284510076, "video_filename": "0025 - Heico - Sigma - 6.avi"} +{"video_id": "0026", "status": "completed", "started_at": "2026-08-25T02:34:45.299569+08:00", "intervals": [{"start_frame": 104407, "end_frame": 104647, "start_seconds": 4176.28, "end_seconds": 4185.92, "duration_seconds": 9.64, "start_time": "01:09:36.280", "end_time": "01:09:45.920", "frame_count": 241, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0026/interval_001_before_f104406.jpg", "start": "evidence/0026/interval_001_start_f104407.jpg", "end": "evidence/0026/interval_001_end_f104647.jpg", "after": "evidence/0026/interval_001_after_f104648.jpg"}}, {"start_frame": 106008, "end_frame": 106558, "start_seconds": 4240.32, "end_seconds": 4262.36, "duration_seconds": 22.04, "start_time": "01:10:40.320", "end_time": "01:11:02.360", "frame_count": 551, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0026/interval_002_before_f106007.jpg", "start": "evidence/0026/interval_002_start_f106008.jpg", "end": "evidence/0026/interval_002_end_f106558.jpg", "after": "evidence/0026/interval_002_after_f106559.jpg"}}, {"start_frame": 135166, "end_frame": 135550, "start_seconds": 5406.64, "end_seconds": 5422.04, "duration_seconds": 15.4, "start_time": "01:30:06.640", "end_time": "01:30:22.040", "frame_count": 385, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0026/interval_003_before_f135165.jpg", "start": "evidence/0026/interval_003_start_f135166.jpg", "end": "evidence/0026/interval_003_end_f135550.jpg", "after": "evidence/0026/interval_003_after_f135551.jpg"}}, {"start_frame": 162447, "end_frame": 162857, "start_seconds": 6497.88, "end_seconds": 6514.32, "duration_seconds": 16.44, "start_time": "01:48:17.880", "end_time": "01:48:34.320", "frame_count": 411, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0026/interval_004_before_f162446.jpg", "start": "evidence/0026/interval_004_start_f162447.jpg", "end": "evidence/0026/interval_004_end_f162857.jpg", "after": "evidence/0026/interval_004_after_f162858.jpg"}}, {"start_frame": 179005, "end_frame": 180062, "start_seconds": 7160.2, "end_seconds": 7202.52, "duration_seconds": 42.32, "start_time": "01:59:20.200", "end_time": "02:00:02.520", "frame_count": 1058, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0026/interval_005_before_f179004.jpg", "start": "evidence/0026/interval_005_start_f179005.jpg", "end": "evidence/0026/interval_005_end_f180062.jpg", "after": "evidence/0026/interval_005_after_f180063.jpg"}}, {"start_frame": 184200, "end_frame": 224505, "start_seconds": 7368.0, "end_seconds": 8980.24, "duration_seconds": 1612.24, "start_time": "02:02:48.000", "end_time": "02:29:40.240", "frame_count": 40306, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0026/interval_006_before_f184199.jpg", "start": "evidence/0026/interval_006_start_f184200.jpg", "end": "evidence/0026/interval_006_end_f224505.jpg", "after": "evidence/0026/interval_006_after_f224506.jpg"}}, {"start_frame": 236957, "end_frame": 238289, "start_seconds": 9478.28, "end_seconds": 9531.6, "duration_seconds": 53.32, "start_time": "02:37:58.280", "end_time": "02:38:51.600", "frame_count": 1333, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0026/interval_007_before_f236956.jpg", "start": "evidence/0026/interval_007_start_f236957.jpg", "end": "evidence/0026/interval_007_end_f238289.jpg", "after": "evidence/0026/interval_007_after_f238290.jpg"}}], "fps": 25.0, "frame_count": 240144, "width": 960, "height": 540, "duration_seconds": 9605.76, "decoded_frames": 240144, "coarse_samples": 9606, "coarse_candidate_count": 7, "interval_count": 7, "blue_screen": true, "blue_screen_seconds": 1771.3999999999999, "blue_screen_ratio": 0.18441018722100072, "finished_at": "2026-08-25T02:39:04.130312+08:00", "elapsed_seconds": 258.8307739328593, "video_filename": "0026 - Heico - Sigma - 7.avi"} +{"video_id": "0027", "status": "completed", "started_at": "2026-08-25T02:39:04.147203+08:00", "intervals": [{"start_frame": 1424, "end_frame": 3592, "start_seconds": 56.96, "end_seconds": 143.72, "duration_seconds": 86.76, "start_time": "00:00:56.960", "end_time": "00:02:23.720", "frame_count": 2169, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_001_before_f1423.jpg", "start": "evidence/0027/interval_001_start_f1424.jpg", "end": "evidence/0027/interval_001_end_f3592.jpg", "after": "evidence/0027/interval_001_after_f3593.jpg"}}, {"start_frame": 5305, "end_frame": 5598, "start_seconds": 212.2, "end_seconds": 223.96, "duration_seconds": 11.76, "start_time": "00:03:32.200", "end_time": "00:03:43.960", "frame_count": 294, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_002_before_f5304.jpg", "start": "evidence/0027/interval_002_start_f5305.jpg", "end": "evidence/0027/interval_002_end_f5598.jpg", "after": "evidence/0027/interval_002_after_f5599.jpg"}}, {"start_frame": 34580, "end_frame": 34741, "start_seconds": 1383.2, "end_seconds": 1389.68, "duration_seconds": 6.48, "start_time": "00:23:03.200", "end_time": "00:23:09.680", "frame_count": 162, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_003_before_f34579.jpg", "start": "evidence/0027/interval_003_start_f34580.jpg", "end": "evidence/0027/interval_003_end_f34741.jpg", "after": "evidence/0027/interval_003_after_f34742.jpg"}}, {"start_frame": 35894, "end_frame": 36166, "start_seconds": 1435.76, "end_seconds": 1446.68, "duration_seconds": 10.92, "start_time": "00:23:55.760", "end_time": "00:24:06.680", "frame_count": 273, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_004_before_f35893.jpg", "start": "evidence/0027/interval_004_start_f35894.jpg", "end": "evidence/0027/interval_004_end_f36166.jpg", "after": "evidence/0027/interval_004_after_f36167.jpg"}}, {"start_frame": 76803, "end_frame": 119032, "start_seconds": 3072.12, "end_seconds": 4761.32, "duration_seconds": 1689.2, "start_time": "00:51:12.120", "end_time": "01:19:21.320", "frame_count": 42230, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_005_before_f76802.jpg", "start": "evidence/0027/interval_005_start_f76803.jpg", "end": "evidence/0027/interval_005_end_f119032.jpg", "after": "evidence/0027/interval_005_after_f119033.jpg"}}, {"start_frame": 136180, "end_frame": 136435, "start_seconds": 5447.2, "end_seconds": 5457.44, "duration_seconds": 10.24, "start_time": "01:30:47.200", "end_time": "01:30:57.440", "frame_count": 256, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_006_before_f136179.jpg", "start": "evidence/0027/interval_006_start_f136180.jpg", "end": "evidence/0027/interval_006_end_f136435.jpg", "after": "evidence/0027/interval_006_after_f136436.jpg"}}, {"start_frame": 136515, "end_frame": 136760, "start_seconds": 5460.6, "end_seconds": 5470.44, "duration_seconds": 9.84, "start_time": "01:31:00.600", "end_time": "01:31:10.440", "frame_count": 246, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_007_before_f136514.jpg", "start": "evidence/0027/interval_007_start_f136515.jpg", "end": "evidence/0027/interval_007_end_f136760.jpg", "after": "evidence/0027/interval_007_after_f136761.jpg"}}, {"start_frame": 136865, "end_frame": 137086, "start_seconds": 5474.6, "end_seconds": 5483.48, "duration_seconds": 8.88, "start_time": "01:31:14.600", "end_time": "01:31:23.480", "frame_count": 222, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_008_before_f136864.jpg", "start": "evidence/0027/interval_008_start_f136865.jpg", "end": "evidence/0027/interval_008_end_f137086.jpg", "after": "evidence/0027/interval_008_after_f137087.jpg"}}, {"start_frame": 138468, "end_frame": 138680, "start_seconds": 5538.72, "end_seconds": 5547.24, "duration_seconds": 8.52, "start_time": "01:32:18.720", "end_time": "01:32:27.240", "frame_count": 213, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_009_before_f138467.jpg", "start": "evidence/0027/interval_009_start_f138468.jpg", "end": "evidence/0027/interval_009_end_f138680.jpg", "after": "evidence/0027/interval_009_after_f138681.jpg"}}, {"start_frame": 154029, "end_frame": 154743, "start_seconds": 6161.16, "end_seconds": 6189.76, "duration_seconds": 28.6, "start_time": "01:42:41.160", "end_time": "01:43:09.760", "frame_count": 715, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_010_before_f154028.jpg", "start": "evidence/0027/interval_010_start_f154029.jpg", "end": "evidence/0027/interval_010_end_f154743.jpg", "after": "evidence/0027/interval_010_after_f154744.jpg"}}, {"start_frame": 155735, "end_frame": 155951, "start_seconds": 6229.4, "end_seconds": 6238.08, "duration_seconds": 8.68, "start_time": "01:43:49.400", "end_time": "01:43:58.080", "frame_count": 217, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_011_before_f155734.jpg", "start": "evidence/0027/interval_011_start_f155735.jpg", "end": "evidence/0027/interval_011_end_f155951.jpg", "after": "evidence/0027/interval_011_after_f155952.jpg"}}, {"start_frame": 170430, "end_frame": 171242, "start_seconds": 6817.2, "end_seconds": 6849.72, "duration_seconds": 32.52, "start_time": "01:53:37.200", "end_time": "01:54:09.720", "frame_count": 813, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_012_before_f170429.jpg", "start": "evidence/0027/interval_012_start_f170430.jpg", "end": "evidence/0027/interval_012_end_f171242.jpg", "after": "evidence/0027/interval_012_after_f171243.jpg"}}, {"start_frame": 171760, "end_frame": 173589, "start_seconds": 6870.4, "end_seconds": 6943.6, "duration_seconds": 73.2, "start_time": "01:54:30.400", "end_time": "01:55:43.600", "frame_count": 1830, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_013_before_f171759.jpg", "start": "evidence/0027/interval_013_start_f171760.jpg", "end": "evidence/0027/interval_013_end_f173589.jpg", "after": "evidence/0027/interval_013_after_f173590.jpg"}}, {"start_frame": 175565, "end_frame": 176292, "start_seconds": 7022.6, "end_seconds": 7051.72, "duration_seconds": 29.12, "start_time": "01:57:02.600", "end_time": "01:57:31.720", "frame_count": 728, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_014_before_f175564.jpg", "start": "evidence/0027/interval_014_start_f175565.jpg", "end": "evidence/0027/interval_014_end_f176292.jpg", "after": "evidence/0027/interval_014_after_f176293.jpg"}}, {"start_frame": 188701, "end_frame": 189011, "start_seconds": 7548.04, "end_seconds": 7560.48, "duration_seconds": 12.44, "start_time": "02:05:48.040", "end_time": "02:06:00.480", "frame_count": 311, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_015_before_f188700.jpg", "start": "evidence/0027/interval_015_start_f188701.jpg", "end": "evidence/0027/interval_015_end_f189011.jpg", "after": "evidence/0027/interval_015_after_f189012.jpg"}}, {"start_frame": 193088, "end_frame": 193299, "start_seconds": 7723.52, "end_seconds": 7732.0, "duration_seconds": 8.48, "start_time": "02:08:43.520", "end_time": "02:08:52.000", "frame_count": 212, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_016_before_f193087.jpg", "start": "evidence/0027/interval_016_start_f193088.jpg", "end": "evidence/0027/interval_016_end_f193299.jpg", "after": "evidence/0027/interval_016_after_f193300.jpg"}}, {"start_frame": 193457, "end_frame": 193676, "start_seconds": 7738.28, "end_seconds": 7747.08, "duration_seconds": 8.8, "start_time": "02:08:58.280", "end_time": "02:09:07.080", "frame_count": 220, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_017_before_f193456.jpg", "start": "evidence/0027/interval_017_start_f193457.jpg", "end": "evidence/0027/interval_017_end_f193676.jpg", "after": "evidence/0027/interval_017_after_f193677.jpg"}}, {"start_frame": 204905, "end_frame": 205339, "start_seconds": 8196.2, "end_seconds": 8213.6, "duration_seconds": 17.4, "start_time": "02:16:36.200", "end_time": "02:16:53.600", "frame_count": 435, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0027/interval_018_before_f204904.jpg", "start": "evidence/0027/interval_018_start_f204905.jpg", "end": "evidence/0027/interval_018_end_f205339.jpg", "after": "evidence/0027/interval_018_after_f205340.jpg"}}], "fps": 25.0, "frame_count": 217434, "width": 960, "height": 540, "duration_seconds": 8697.36, "decoded_frames": 217434, "coarse_samples": 8698, "coarse_candidate_count": 18, "interval_count": 18, "blue_screen": true, "blue_screen_seconds": 2061.84, "blue_screen_ratio": 0.23706504042606033, "finished_at": "2026-08-25T02:42:44.422030+08:00", "elapsed_seconds": 220.27485602535307, "video_filename": "0027 - Heico - Sigma - 8.avi"} +{"video_id": "0028", "status": "completed", "started_at": "2026-08-25T02:42:44.441032+08:00", "intervals": [{"start_frame": 2902, "end_frame": 4698, "start_seconds": 116.08, "end_seconds": 187.96, "duration_seconds": 71.88, "start_time": "00:01:56.080", "end_time": "00:03:07.960", "frame_count": 1797, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_001_before_f2901.jpg", "start": "evidence/0028/interval_001_start_f2902.jpg", "end": "evidence/0028/interval_001_end_f4698.jpg", "after": "evidence/0028/interval_001_after_f4699.jpg"}}, {"start_frame": 4812, "end_frame": 5000, "start_seconds": 192.48, "end_seconds": 200.04, "duration_seconds": 7.56, "start_time": "00:03:12.480", "end_time": "00:03:20.040", "frame_count": 189, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_002_before_f4811.jpg", "start": "evidence/0028/interval_002_start_f4812.jpg", "end": "evidence/0028/interval_002_end_f5000.jpg", "after": "evidence/0028/interval_002_after_f5001.jpg"}}, {"start_frame": 5332, "end_frame": 5945, "start_seconds": 213.28, "end_seconds": 237.84, "duration_seconds": 24.56, "start_time": "00:03:33.280", "end_time": "00:03:57.840", "frame_count": 614, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_003_before_f5331.jpg", "start": "evidence/0028/interval_003_start_f5332.jpg", "end": "evidence/0028/interval_003_end_f5945.jpg", "after": "evidence/0028/interval_003_after_f5946.jpg"}}, {"start_frame": 6263, "end_frame": 6645, "start_seconds": 250.52, "end_seconds": 265.84, "duration_seconds": 15.32, "start_time": "00:04:10.520", "end_time": "00:04:25.840", "frame_count": 383, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_004_before_f6262.jpg", "start": "evidence/0028/interval_004_start_f6263.jpg", "end": "evidence/0028/interval_004_end_f6645.jpg", "after": "evidence/0028/interval_004_after_f6646.jpg"}}, {"start_frame": 81620, "end_frame": 82270, "start_seconds": 3264.8, "end_seconds": 3290.84, "duration_seconds": 26.04, "start_time": "00:54:24.800", "end_time": "00:54:50.840", "frame_count": 651, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_005_before_f81619.jpg", "start": "evidence/0028/interval_005_start_f81620.jpg", "end": "evidence/0028/interval_005_end_f82270.jpg", "after": "evidence/0028/interval_005_after_f82271.jpg"}}, {"start_frame": 83668, "end_frame": 84298, "start_seconds": 3346.72, "end_seconds": 3371.96, "duration_seconds": 25.24, "start_time": "00:55:46.720", "end_time": "00:56:11.960", "frame_count": 631, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_006_before_f83667.jpg", "start": "evidence/0028/interval_006_start_f83668.jpg", "end": "evidence/0028/interval_006_end_f84298.jpg", "after": "evidence/0028/interval_006_after_f84299.jpg"}}, {"start_frame": 91175, "end_frame": 92008, "start_seconds": 3647.0, "end_seconds": 3680.36, "duration_seconds": 33.36, "start_time": "01:00:47.000", "end_time": "01:01:20.360", "frame_count": 834, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_007_before_f91174.jpg", "start": "evidence/0028/interval_007_start_f91175.jpg", "end": "evidence/0028/interval_007_end_f92008.jpg", "after": "evidence/0028/interval_007_after_f92009.jpg"}}, {"start_frame": 98941, "end_frame": 99826, "start_seconds": 3957.64, "end_seconds": 3993.08, "duration_seconds": 35.44, "start_time": "01:05:57.640", "end_time": "01:06:33.080", "frame_count": 886, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_008_before_f98940.jpg", "start": "evidence/0028/interval_008_start_f98941.jpg", "end": "evidence/0028/interval_008_end_f99826.jpg", "after": "evidence/0028/interval_008_after_f99827.jpg"}}, {"start_frame": 99972, "end_frame": 100157, "start_seconds": 3998.88, "end_seconds": 4006.32, "duration_seconds": 7.44, "start_time": "01:06:38.880", "end_time": "01:06:46.320", "frame_count": 186, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_009_before_f99971.jpg", "start": "evidence/0028/interval_009_start_f99972.jpg", "end": "evidence/0028/interval_009_end_f100157.jpg", "after": "evidence/0028/interval_009_after_f100158.jpg"}}, {"start_frame": 106430, "end_frame": 107644, "start_seconds": 4257.2, "end_seconds": 4305.8, "duration_seconds": 48.6, "start_time": "01:10:57.200", "end_time": "01:11:45.800", "frame_count": 1215, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_010_before_f106429.jpg", "start": "evidence/0028/interval_010_start_f106430.jpg", "end": "evidence/0028/interval_010_end_f107644.jpg", "after": "evidence/0028/interval_010_after_f107645.jpg"}}, {"start_frame": 108850, "end_frame": 109555, "start_seconds": 4354.0, "end_seconds": 4382.24, "duration_seconds": 28.24, "start_time": "01:12:34.000", "end_time": "01:13:02.240", "frame_count": 706, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_011_before_f108849.jpg", "start": "evidence/0028/interval_011_start_f108850.jpg", "end": "evidence/0028/interval_011_end_f109555.jpg", "after": "evidence/0028/interval_011_after_f109556.jpg"}}, {"start_frame": 111863, "end_frame": 113239, "start_seconds": 4474.52, "end_seconds": 4529.6, "duration_seconds": 55.08, "start_time": "01:14:34.520", "end_time": "01:15:29.600", "frame_count": 1377, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_012_before_f111862.jpg", "start": "evidence/0028/interval_012_start_f111863.jpg", "end": "evidence/0028/interval_012_end_f113239.jpg", "after": "evidence/0028/interval_012_after_f113240.jpg"}}, {"start_frame": 113515, "end_frame": 114657, "start_seconds": 4540.6, "end_seconds": 4586.32, "duration_seconds": 45.72, "start_time": "01:15:40.600", "end_time": "01:16:26.320", "frame_count": 1143, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_013_before_f113514.jpg", "start": "evidence/0028/interval_013_start_f113515.jpg", "end": "evidence/0028/interval_013_end_f114657.jpg", "after": "evidence/0028/interval_013_after_f114658.jpg"}}, {"start_frame": 129364, "end_frame": 146585, "start_seconds": 5174.56, "end_seconds": 5863.44, "duration_seconds": 688.88, "start_time": "01:26:14.560", "end_time": "01:37:43.440", "frame_count": 17222, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_014_before_f129363.jpg", "start": "evidence/0028/interval_014_start_f129364.jpg", "end": "evidence/0028/interval_014_end_f146585.jpg", "after": "evidence/0028/interval_014_after_f146586.jpg"}}, {"start_frame": 157535, "end_frame": 157730, "start_seconds": 6301.4, "end_seconds": 6309.24, "duration_seconds": 7.84, "start_time": "01:45:01.400", "end_time": "01:45:09.240", "frame_count": 196, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_015_before_f157534.jpg", "start": "evidence/0028/interval_015_start_f157535.jpg", "end": "evidence/0028/interval_015_end_f157730.jpg", "after": "evidence/0028/interval_015_after_f157731.jpg"}}, {"start_frame": 164683, "end_frame": 165551, "start_seconds": 6587.32, "end_seconds": 6622.08, "duration_seconds": 34.76, "start_time": "01:49:47.320", "end_time": "01:50:22.080", "frame_count": 869, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_016_before_f164682.jpg", "start": "evidence/0028/interval_016_start_f164683.jpg", "end": "evidence/0028/interval_016_end_f165551.jpg", "after": "evidence/0028/interval_016_after_f165552.jpg"}}, {"start_frame": 178354, "end_frame": 178553, "start_seconds": 7134.16, "end_seconds": 7142.16, "duration_seconds": 8.0, "start_time": "01:58:54.160", "end_time": "01:59:02.160", "frame_count": 200, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_017_before_f178353.jpg", "start": "evidence/0028/interval_017_start_f178354.jpg", "end": "evidence/0028/interval_017_end_f178553.jpg", "after": "evidence/0028/interval_017_after_f178554.jpg"}}, {"start_frame": 183895, "end_frame": 184094, "start_seconds": 7355.8, "end_seconds": 7363.8, "duration_seconds": 8.0, "start_time": "02:02:35.800", "end_time": "02:02:43.800", "frame_count": 200, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_018_before_f183894.jpg", "start": "evidence/0028/interval_018_start_f183895.jpg", "end": "evidence/0028/interval_018_end_f184094.jpg", "after": "evidence/0028/interval_018_after_f184095.jpg"}}, {"start_frame": 188618, "end_frame": 234529, "start_seconds": 7544.72, "end_seconds": 9381.2, "duration_seconds": 1836.48, "start_time": "02:05:44.720", "end_time": "02:36:21.200", "frame_count": 45912, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_019_before_f188617.jpg", "start": "evidence/0028/interval_019_start_f188618.jpg", "end": "evidence/0028/interval_019_end_f234529.jpg", "after": "evidence/0028/interval_019_after_f234530.jpg"}}, {"start_frame": 284965, "end_frame": 286101, "start_seconds": 11398.6, "end_seconds": 11444.08, "duration_seconds": 45.48, "start_time": "03:09:58.600", "end_time": "03:10:44.080", "frame_count": 1137, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_020_before_f284964.jpg", "start": "evidence/0028/interval_020_start_f284965.jpg", "end": "evidence/0028/interval_020_end_f286101.jpg", "after": "evidence/0028/interval_020_after_f286102.jpg"}}, {"start_frame": 315904, "end_frame": 316225, "start_seconds": 12636.16, "end_seconds": 12649.04, "duration_seconds": 12.88, "start_time": "03:30:36.160", "end_time": "03:30:49.040", "frame_count": 322, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_021_before_f315903.jpg", "start": "evidence/0028/interval_021_start_f315904.jpg", "end": "evidence/0028/interval_021_end_f316225.jpg", "after": "evidence/0028/interval_021_after_f316226.jpg"}}, {"start_frame": 319567, "end_frame": 364856, "start_seconds": 12782.68, "end_seconds": 14594.28, "duration_seconds": 1811.6, "start_time": "03:33:02.680", "end_time": "04:03:14.280", "frame_count": 45290, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_022_before_f319566.jpg", "start": "evidence/0028/interval_022_start_f319567.jpg", "end": "evidence/0028/interval_022_end_f364856.jpg", "after": "evidence/0028/interval_022_after_f364857.jpg"}}, {"start_frame": 381302, "end_frame": 381504, "start_seconds": 15252.08, "end_seconds": 15260.2, "duration_seconds": 8.12, "start_time": "04:14:12.080", "end_time": "04:14:20.200", "frame_count": 203, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0028/interval_023_before_f381301.jpg", "start": "evidence/0028/interval_023_start_f381302.jpg", "end": "evidence/0028/interval_023_end_f381504.jpg", "after": "evidence/0028/interval_023_after_f381505.jpg"}}], "fps": 25.0, "frame_count": 403110, "width": 960, "height": 540, "duration_seconds": 16124.4, "decoded_frames": 403110, "coarse_samples": 16125, "coarse_candidate_count": 24, "interval_count": 23, "blue_screen": true, "blue_screen_seconds": 4886.5199999999995, "blue_screen_ratio": 0.303051276326561, "finished_at": "2026-08-25T02:50:14.300589+08:00", "elapsed_seconds": 449.8595881573856, "video_filename": "0028 - Heico - Sigma - 9.avi"} +{"video_id": "0029", "status": "completed", "started_at": "2026-08-25T02:50:14.346982+08:00", "intervals": [{"start_frame": 446, "end_frame": 620, "start_seconds": 17.84, "end_seconds": 24.84, "duration_seconds": 7.0, "start_time": "00:00:17.840", "end_time": "00:00:24.840", "frame_count": 175, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0029/interval_001_before_f445.jpg", "start": "evidence/0029/interval_001_start_f446.jpg", "end": "evidence/0029/interval_001_end_f620.jpg", "after": "evidence/0029/interval_001_after_f621.jpg"}}, {"start_frame": 1218, "end_frame": 1586, "start_seconds": 48.72, "end_seconds": 63.48, "duration_seconds": 14.76, "start_time": "00:00:48.720", "end_time": "00:01:03.480", "frame_count": 369, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0029/interval_002_before_f1217.jpg", "start": "evidence/0029/interval_002_start_f1218.jpg", "end": "evidence/0029/interval_002_end_f1586.jpg", "after": "evidence/0029/interval_002_after_f1587.jpg"}}, {"start_frame": 168562, "end_frame": 217347, "start_seconds": 6742.48, "end_seconds": 8693.92, "duration_seconds": 1951.44, "start_time": "01:52:22.480", "end_time": "02:24:53.920", "frame_count": 48786, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0029/interval_003_before_f168561.jpg", "start": "evidence/0029/interval_003_start_f168562.jpg", "end": "evidence/0029/interval_003_end_f217347.jpg", "after": "evidence/0029/interval_003_after_f217348.jpg"}}, {"start_frame": 222032, "end_frame": 222368, "start_seconds": 8881.28, "end_seconds": 8894.76, "duration_seconds": 13.48, "start_time": "02:28:01.280", "end_time": "02:28:14.760", "frame_count": 337, "mean_blue_ratio": 1.0, "exact_repeated_hash": "00c9758b75e5aaadf2d346b15f90db76", "confidence": "high", "evidence_frames": {"before": "evidence/0029/interval_004_before_f222031.jpg", "start": "evidence/0029/interval_004_start_f222032.jpg", "end": "evidence/0029/interval_004_end_f222368.jpg", "after": "evidence/0029/interval_004_after_f222369.jpg"}}], "fps": 25.0, "frame_count": 235158, "width": 960, "height": 540, "duration_seconds": 9406.32, "decoded_frames": 235158, "coarse_samples": 9407, "coarse_candidate_count": 4, "interval_count": 4, "blue_screen": true, "blue_screen_seconds": 1986.68, "blue_screen_ratio": 0.2112069332108625, "finished_at": "2026-08-25T02:54:37.053506+08:00", "elapsed_seconds": 262.7065535802394, "video_filename": "0029 - Heico - Sigma - 10.avi"} diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/.dockerignore b/pipelines/Tr006_V2/docker/segment-algorithm/.dockerignore new file mode 100644 index 0000000000000000000000000000000000000000..bf931adf6f183264d57d7ffd1664223f4344740a --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/.dockerignore @@ -0,0 +1,6 @@ +# Keep the build context small. do_save.sh writes multi-gigabyte *.tar.gz +# archives into this folder, and without this file Docker would send them +# (and the sample batch) to the daemon on every build. +*.tar.gz +test/ +__pycache__/ diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/.gitignore b/pipelines/Tr006_V2/docker/segment-algorithm/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..489d05f3aa750c70916c785ba30f9b9647af1269 --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/.gitignore @@ -0,0 +1,2 @@ +*.tar.gz +__pycache__/ diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/Dockerfile b/pipelines/Tr006_V2/docker/segment-algorithm/Dockerfile new file mode 100644 index 0000000000000000000000000000000000000000..6c14b6f8a37567dcbf9bb10a6ed144d22f05a019 --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/Dockerfile @@ -0,0 +1,102 @@ +FROM --platform=linux/amd64 pytorch/pytorch:2.11.0-cuda12.8-cudnn9-runtime AS segment-algorithm-amd64 + +# torch 2.11.0 + torchvision are pre-installed in the base image (CUDA 12.8). +# requirements.txt only lists additional dependencies. +# +# This base image is the recommended default, not a hard requirement. It is chosen +# because it is the ONE PyTorch build that runs on both evaluation GPUs unmodified: +# an NVIDIA RTX PRO 6000 Blackwell (compute capability sm_120) and an NVIDIA L40S +# (sm_89). Its CUDA 12.8 kernels cover sm_75/80/86/90/100/120 — Blackwell natively +# and the L40S via its sm_86 kernels (a CUDA binary runs on any later minor +# revision of the same major capability). So one image serves whichever GPU you +# select on the platform, with no rebuild. +# +# If you change it, know the two limits it steers around: CUDA 12.4/12.6 wheels +# stop at sm_90 and abort on Blackwell ("not compatible with the current PyTorch +# installation"); CUDA 13.0 wheels cover both GPUs but need a host driver >= 580 — +# the RTX PRO 6000 fleet has it (595), the L40S fleet does not. If your submission +# targets ONLY the RTX PRO 6000 (offered on the SEGMENT and PROCEDURE tracks), a +# newer CUDA/torch is a valid choice — see the guard below for the one line to relax. + +ENV PYTHONUNBUFFERED=1 + +# ── System packages ─────────────────────────────────────────────────────────── +# The slim -runtime base image ships neither the shared libraries that +# opencv-python (a dependency of orena-focus) needs at import time, nor any +# compiler toolchain. Install system packages HERE, while the build still runs +# as root — everything below runs as the unprivileged 'user'. If one of your +# pip dependencies compiles from source (e.g. flash-attn), add build-essential. +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + libgl1 \ + libglib2.0-0 \ + libxcb1 \ + && rm -rf /var/lib/apt/lists/* + +RUN groupadd -r user && useradd -m --no-log-init -r -g user user +USER user + +WORKDIR /opt/app + +COPY --chown=user:user requirements.txt /opt/app/ + +# If one of your dependencies compiles CUDA extensions during the build (e.g. +# flash-attn from source), uncomment this so it emits kernels for both evaluation +# GPUs rather than only for the card in your build machine: +# ENV TORCH_CUDA_ARCH_LIST="8.0;8.6;8.9;9.0;12.0+PTX" + +# --break-system-packages: this base image's Python is Ubuntu's system Python, +# which is marked externally-managed (PEP 668), so pip refuses to install without +# it. Nothing is actually broken — `--user` still installs into the unprivileged +# user's own ~/.local, leaving the system packages untouched. +RUN python -m pip install \ + --user \ + --break-system-packages \ + --no-cache-dir \ + --no-color \ + --requirement /opt/app/requirements.txt + +# ── Guard: is torch still the recommended CUDA 12.8 build? ──────────────────── +# This enforces the recommendation above so an ACCIDENTAL swap cannot reach the +# platform unnoticed. If one of your dependencies needs a torch version this image +# does not have, pip quietly replaces the CUDA 12.8 build with the default PyPI +# wheel (built against CUDA 13.0). That wheel needs a host driver >= 580.65.06; the +# L40S fleet runs 570.211.01, so it cannot start there. Failing the build now beats +# failing every job after you submit. +# +# Both halves matter. The CUDA version is what the driver floor depends on, and it +# is NOT visible in the architecture list: the CUDA 13.0 wheel advertises the very +# same sm_ list as the 12.8 one. The arch list is still checked because it is what +# an older CUDA 12.4/12.6 wheel gets wrong. get_arch_list() returns [] with no GPU +# visible — a docker build never has one — so the flags compiled into the wheel are +# read directly. +# +# Targeting the RTX PRO 6000 only, and want a newer CUDA? Then this default no +# longer applies to you: drop the `cuda.startswith('12.8')` assert below (keep the +# sm_120 arch assert, so a build with no Blackwell kernels is still caught). +RUN python -c "import torch; \ +flags = torch._C._cuda_getArchFlags() or ''; \ +cuda = torch.version.cuda or ''; \ +print('torch', torch.__version__, '| cuda', cuda, '|', flags); \ +assert cuda.startswith('12.8'), 'torch was replaced by a CUDA ' + cuda + ' build; the L40S driver cannot run it'; \ +assert 'sm_120' in flags and 'sm_86' in flags, 'torch build misses a required GPU architecture: ' + repr(flags)" + +# ── Model definition + weights ──────────────────────────────────────────────── +# resources/ holds the model definition (model.py) and a small committed dummy +# checkpoint (dummy_weights.pt). For a real submission, replace dummy_weights.pt +# with your trained checkpoint. +# +# ECR layer-size limit: a single image layer must not exceed 50 GB, and each +# COPY instruction produces exactly one layer. A checkpoint large enough to push +# one COPY past that limit must be split into chunks and copied with several COPY +# instructions (one layer each), then reassembled at runtime — e.g. pre-split +# with `split -b 45G weights.pt resources/weights.part-` and: +# COPY --chown=user:user resources/weights.part-aa /opt/app/resources/ +# COPY --chown=user:user resources/weights.part-ab /opt/app/resources/ +# COPY --chown=user:user resources/weights.part-ac /opt/app/resources/ +# reassembling with `cat resources/weights.part-* > weights.pt` before load. +COPY --chown=user:user resources/ /opt/app/resources/ + +COPY --chown=user:user inference.py /opt/app/ + +ENTRYPOINT ["python", "inference.py"] diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/do_build.sh b/pipelines/Tr006_V2/docker/segment-algorithm/do_build.sh new file mode 100644 index 0000000000000000000000000000000000000000..aec1b5536efb7b07862b9f1d7e8f225817110e45 --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/do_build.sh @@ -0,0 +1,11 @@ +#!/usr/bin/env bash + +set -e + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +DOCKER_IMAGE_TAG="segment-algorithm" + +docker build \ + --platform=linux/amd64 \ + --tag "$DOCKER_IMAGE_TAG" \ + "$SCRIPT_DIR" 2>&1 diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/do_save.sh b/pipelines/Tr006_V2/docker/segment-algorithm/do_save.sh new file mode 100644 index 0000000000000000000000000000000000000000..fd676c05777e723ad719a36450d874bdfcd03221 --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/do_save.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash + +set -eo pipefail + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) + +DOCKER_IMAGE_TAG="segment-algorithm" + +echo "=+= (Re)build the container" +source "${SCRIPT_DIR}/do_build.sh" + +build_timestamp=$( docker inspect --format='{{ .Created }}' "$DOCKER_IMAGE_TAG") + +if [ -z "$build_timestamp" ]; then + echo "Error: Failed to retrieve build information for container $DOCKER_IMAGE_TAG" + exit 1 +fi + +formatted_build_info=$(echo $build_timestamp | sed -E 's/(.*)T(.*)\..*Z/\1_\2/' | sed 's/[-,:]/-/g') + +output_filename="${SCRIPT_DIR}/${DOCKER_IMAGE_TAG}_${formatted_build_info}.tar.gz" + +echo "==+==" +echo "Saving the container image as ${output_filename}. This can take a while." +echo "" + +# Use pigz (parallel gzip) if available — much faster than gzip on a large image. +if command -v pigz >/dev/null 2>&1; then + compressor=(pigz -c) +else + compressor=(gzip -c) +fi + +# Show a progress bar with pv if available, sized by the image's uncompressed size. +image_size=$(docker image inspect "$DOCKER_IMAGE_TAG" --format '{{.Size}}' 2>/dev/null || true) + +if command -v pv >/dev/null 2>&1; then + if [[ "$image_size" =~ ^[0-9]+$ && "$image_size" -gt 0 ]]; then + docker save "$DOCKER_IMAGE_TAG" \ + | pv --size "$image_size" --progress --timer --eta --rate --bytes \ + | "${compressor[@]}" > "$output_filename" + else + docker save "$DOCKER_IMAGE_TAG" | pv --timer --rate --bytes | "${compressor[@]}" > "$output_filename" + fi +else + echo "Tip: install 'pv' (e.g. sudo apt-get install pv) for a progress bar." + docker save "$DOCKER_IMAGE_TAG" | "${compressor[@]}" > "$output_filename" +fi +echo "Container image saved as ${output_filename}" +echo "==+==" diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/do_test_run.sh b/pipelines/Tr006_V2/docker/segment-algorithm/do_test_run.sh new file mode 100644 index 0000000000000000000000000000000000000000..fd61ef9a51722da1583da8ba025aa57dcfc1423d --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/do_test_run.sh @@ -0,0 +1,86 @@ +#!/usr/bin/env bash + +set -e + +SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd ) +DOCKER_IMAGE_TAG="segment-algorithm" + +DOCKER_NOOP_VOLUME="${DOCKER_IMAGE_TAG}-volume" + +INPUT_DIR="${SCRIPT_DIR}/test/input" +OUTPUT_DIR="${SCRIPT_DIR}/test/output" + +# Use the GPU when Docker's NVIDIA runtime is available; fall back to CPU +# otherwise (the dummy runs fine on CPU — the platform always provides a GPU). +GPU_FLAGS=() +if docker info --format '{{range $k, $v := .Runtimes}}{{$k}} {{end}}' 2>/dev/null | grep -qw nvidia; then + GPU_FLAGS=(--gpus all) +else + echo "=+= No NVIDIA container runtime detected - running on CPU" +fi + +echo "=+= (Re)build the container" +source "${SCRIPT_DIR}/do_build.sh" + +cleanup() { + echo "=+= Cleaning permissions ..." + docker run --rm \ + --platform=linux/amd64 \ + --quiet \ + --volume "$OUTPUT_DIR":/output \ + --entrypoint /bin/sh \ + $DOCKER_IMAGE_TAG \ + -c "chmod -R -f o+rwX /output/* || true" + + docker volume rm "$DOCKER_NOOP_VOLUME" > /dev/null +} + +chmod -R -f o+rX "$INPUT_DIR" + +if [ -d "${OUTPUT_DIR}/interface_1" ]; then + chmod -f o+rwX "${OUTPUT_DIR}/interface_1" + + echo "=+= Cleaning up any earlier output" + docker run --rm \ + --platform=linux/amd64 \ + --quiet \ + --volume "${OUTPUT_DIR}/interface_1":/output \ + --entrypoint /bin/sh \ + $DOCKER_IMAGE_TAG \ + -c "rm -rf /output/* || true" +else + mkdir -p -m o+rwX "${OUTPUT_DIR}/interface_1" +fi + +docker volume create "$DOCKER_NOOP_VOLUME" > /dev/null + +trap cleanup EXIT + +run_docker_forward_pass() { + local interface_dir="$1" + + echo "=+= Doing a forward pass on ${interface_dir}" + + ## Note the extra arguments that are passed here: + # '--network none' + # entails there is no internet connection + # '--volume :/tmp' + # is added because on the ORena FOCUS Challenge platform this directory cannot be used to store permanent files + # '--shm-size' is left at Docker's default (64 MB). A PyTorch DataLoader with + # num_workers > 0 needs more shared memory — if you use one, add e.g. + # '--shm-size=2g' here. + docker run --rm \ + --platform=linux/amd64 \ + --network none \ + "${GPU_FLAGS[@]}" \ + --volume "${INPUT_DIR}/${interface_dir}":/input:ro \ + --volume "${OUTPUT_DIR}/${interface_dir}":/output \ + --volume "$DOCKER_NOOP_VOLUME":/tmp \ + "$DOCKER_IMAGE_TAG" + + echo "=+= Wrote results to ${OUTPUT_DIR}/${interface_dir}" +} + +run_docker_forward_pass "interface_1" + +echo "=+= Save this image for uploading via ./do_save.sh" diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/inference.py b/pipelines/Tr006_V2/docker/segment-algorithm/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..12ba2d37221162d942c036815f814bc4fcbf76d4 --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/inference.py @@ -0,0 +1,393 @@ +""" +Dummy algorithm for the ORena SAVE FOCUS challenge — SEGMENT track. + +Replace the marked blocks with your model. Every spot to change carries a ``TODO``. + +What one run receives +--------------------- +A single container run answers a **batch of questions**. +Inputs are mounted read-only at ``/input``: + + request.json a LIST of focus.Request objects — one per question + plain/.mp4 the video clip for question + overlayed/.mp4 the same clip with an HH:MM:SS timestamp burned in + FO_definitions.json the foreign-object class definitions, as plain text + batch.json a convenience index of the batch (qIDs and layout) + +and one file is written to ``/output``: + + answer.json a LIST of focus.Response objects — one per question + +``request.json`` is authoritative: iterate it, and answer every ``qID`` it lists. + +Each question carries its own clip +---------------------------------- +Every question has a **separate video file, already trimmed to that question's +window**, and no two questions in a batch come from the same source video. Your +answer to one question can therefore never depend on another question's footage. + +Two consequences for your pipeline: + +* ``request.start_time`` / ``end_time`` refer to the **original, untrimmed + procedure video** — the same timeline shown in the overlay variant. The clip you + are given has already been cut to that window, so decode the clip from its own + beginning; do **not** seek to ``start_time`` inside it. ``request.duration`` + still tells you how long the window is. +* There is nothing to share between questions — no common video to cache or encode + once. Each clip is decoded on its own. + +Both variants show the same footage. Use the plain clip, the overlay, or both. + +Clip encoding +------------- +Clips are H.264 MP4 at exactly **5 fps** (one frame every 0.2 s), height-normalised +to at most 576 px. Widths follow the source's aspect ratio (e.g., 1024x576, 720x576 or +640x360) — do not assume a fixed resolution. A +keyframe every 5 s keeps seeking cheap. The overlay clock shows the source +timeline's second, floored. + +Load your model once +-------------------- +Model loading, imports and warm-up happen **once per container run**, not once per +question. The latency budget is built around that, so keep model loading outside +the per-question loop — as this file does. This is the single most important thing +to get right for speed. + +The latency budget +------------------ +The platform measures the execution time of the whole run itself; the ``latency`` +field you write into each Response is informational and is not used for scoring. + +The budget is pooled across the batch. For a batch of B questions: + + allowed = 120 s + B x 15 s + ^setup ^per-question processing (the SEGMENT budget) + +Going over does not fail the run outright. Instead, a share of your questions is +**forfeited** (scored incorrect), growing with the overrun: a small overrun costs a +few questions, and exceeding the budget by 20% forfeits the whole batch. + +The forfeited questions are selected for you — not necessarily the slow ones — so an +overrun costs you questions you would have answered correctly. If you are running out +of budget, answer the rest quickly rather than overrun. + +Keep going when a question fails +-------------------------------- +One bad question should not cost you the batch. Catch errors per question, still +emit a Response for that ``qID``, and carry on — as the loop below does. A question +with no answer is simply scored incorrect. + +Frame sampling +-------------- +The clip is decoded with `decord `_ using a single +batched read of exactly 80 evenly spaced frame indices. The pixels come from the +trimmed clip, while their timestamps are reconstructed on the original timeline +from ``Request.start_time``. Decoding and processor work count towards latency. + +Trying it out +------------- +A sample batch (``test/input/interface_1/``) ships with the template. This package +stores the merged Tr006 checkpoint under ``resources/model``; runtime inference is +fully local and does not download model files. + + ./do_test_run.sh # build the image and run it on the sample batch + ./do_save.sh # build and save a .tar.gz for upload +""" + +import json +import logging +import sys +import time +from pathlib import Path + +import numpy as np +import torch + +# decord must be imported AFTER torch: importing it first breaks CUDA +# initialisation (a known decord issue — torch.cuda init then fails with +# "random_device could not be read" or a segfault). +import decord +from focus import Request, Response, load_requests, save_items +from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration + +logging.basicConfig( + stream=sys.stdout, + level=logging.DEBUG, + format="%(asctime)s [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", +) +log = logging.getLogger(__name__) + +RESOURCES_PATH = Path(__file__).parent / "resources" +MODEL_PATH = RESOURCES_PATH / "model" + +INPUT_PATH = Path("/input") +OUTPUT_PATH = Path("/output") +N_FRAMES = 80 +MAX_NEW_TOKENS = 64 +USE_OVERLAY = False # The trained Tr006 V2 model used plain frames + metadata. +VIDEO_DIR = INPUT_PATH / ("overlayed" if USE_OVERLAY else "plain") + +# Keep this prompt byte-for-byte aligned with the Tr006 V2 fixed80 training +# manifest. The official package supplies FO_definitions separately; it is +# deliberately not concatenated to the model prompt because the training rows +# used this system message and the question only. +SYSTEM_PROMPT = ( + "You are a surgical video analysis assistant. You are shown a short clip from a laparoscopic procedure and asked a single question about it.\n\n" + "The clip is given as a sequence of frames in chronological order. Each frame is preceded by its timestamp in the source video, written as " + " -- a frame marked <1234.5 seconds> was taken 1234.5 seconds (00:20:34) after the video began. Consecutive frames may be several seconds " + "apart, so an event can happen between two of them.\n\n" + "A foreign object (FO) is any object fully introduced into the patient's body\n" + "cavity during surgery that must be retrieved or accounted for. Importantly,\n" + "standard surgical instruments that remain connected to the external environment\n" + "(e.g., graspers, scissors, trocars, staplers, cameras) are not considered foreign\n" + "objects. Furthermore, we exclude detachable parts of surgical instruments,\n" + "particularly anvil components of staplers.\n\n" + "The foreign object classes are exactly: Sponge, Clip, Specimen Bag, Silicone Loop, External Drain, Needle, Gallstone, Specimen, Mesh, Absorbable Hemostatic Agent.\n\n" + "Reply with the answer and nothing else -- no reasoning, no preamble, no\n" + "explanation, no restating the question. A single short line.\n\n" + "Rules for the answer:\n" + "- Write the value only. No sentence, no explanation, no units, no trailing\n" + " period, and never repeat the question.\n" + "- Asks yes or no -> write exactly: yes or no\n" + "- Asks how many / for a count -> write digits only, e.g. 0 or 1 or 2.\n" + "- Asks which foreign object class(es) -> write class names exactly as spelled\n" + " in the list above, comma-separated (e.g. Clip, Sponge), or exactly: none\n" + " Never answer with a generic description such as \"surgical instrument\".\n" + "- Asks when something happens -> write hh:mm:ss, read off the frame timestamps.\n" + " Those timestamps count from the START OF THE VIDEO, not from the start of this\n" + " clip, so <1234.5 seconds> is 00:20:34. Convert carefully: 3600 seconds is one\n" + " hour, 60 seconds is one minute. If the event falls between two frames, give\n" + " your best estimate between their timestamps rather than snapping to one.\n" + " For several time points, separate them with commas: 00:09:19, 00:12:44\n" + "- Asks for a percentage or a share -> write the number only, e.g. 40\n" + "- Lists options to choose from -> copy exactly one of those options, verbatim.\n" + "- Anything else -> a short phrase, at most a few words.\n\n" + "If you are unsure, still commit to your single best answer in the required\n" + "form. An empty, hedged, or explanatory answer is scored as wrong.\n" +) + + +def log_environment(device: torch.device) -> None: + """Report the CUDA stack this container ended up with — once, at startup. + + A torch build carrying no kernels for the GPU it was handed is the classic + reason a run dies on the platform, and it stays invisible until the first + CUDA call. Printing this makes a failed job diagnosable from the logs alone. + Keep it: it costs milliseconds, once per run. + """ + log.info("--- Environment ---") + log.info(" torch : %s (built against CUDA %s)", torch.__version__, torch.version.cuda) + + # get_arch_list() is empty when no GPU is visible; the architecture flags + # compiled into the wheel can be read either way. + arch = torch.cuda.get_arch_list() or (torch._C._cuda_getArchFlags() or "").split() + log.info(" torch kernels : %s", " ".join(arch) or "unknown") + + driver_file = Path("/proc/driver/nvidia/version") + log.info( + " host driver : %s", + driver_file.read_text().strip().splitlines()[0] + if driver_file.exists() + else "no NVIDIA driver visible", + ) + + if device.type != "cuda": + log.warning(" No GPU visible — running on CPU. The platform always provides one.") + return + + capability = torch.cuda.get_device_capability(0) + free, total = torch.cuda.mem_get_info(0) + log.info(" GPU : %s", torch.cuda.get_device_name(0)) + log.info(" capability : %s (sm_%d%d)", capability, *capability) + log.info(" VRAM : %.1f GiB free of %.1f GiB", free / 1024**3, total / 1024**3) + + +def clip_path_for(req: Request) -> Path: + """The video clip belonging to this question. + + Every question has exactly one clip, named after its ``qID`` and already cut to + that question's ``[start_time, end_time]`` window. + """ + return VIDEO_DIR / f"{req.qID}.mp4" + + +def your_preprocessing(video_path: Path, req: Request) -> tuple[np.ndarray, dict]: + """Decode exactly ``N_FRAMES`` and attach original-timeline millisecond PTS. + + Official clips are already trimmed and encoded at 5fps. The local frame + indices are therefore sampled across the whole clip, while ``frames_indices`` + are rebuilt from Request.start_time so Qwen's processor emits absolute + ``<... seconds>`` anchors, matching the Tr006 fixed80 training contract. + """ + log.debug("Decoding %s (size on disk: %d bytes)", video_path, video_path.stat().st_size) + + vr = decord.VideoReader(str(video_path), ctx=decord.cpu(0), num_threads=1) + base_fps = float(vr.get_avg_fps()) + n_frames = len(vr) + if n_frames <= 0 or base_fps <= 0: + raise RuntimeError(f"invalid video stream: n_frames={n_frames}, fps={base_fps}") + # np.linspace + round mirrors the fixed-frame preprocessing and supports + # short clips by repeating the nearest decoded frame. + indices = np.linspace(0, n_frames - 1, N_FRAMES).round().astype(np.int64) + log.debug("Video: %d frames @ %.3f fps; sampling %d across %.2f s", n_frames, base_fps, N_FRAMES, n_frames / base_fps) + frames = vr.get_batch(indices).asnumpy() # (T, H, W, 3) uint8 RGB — one batched decode + del vr + + # Qwen3's processor averages timestamps per temporal patch (2 frames), so + # metadata must have one entry per decoded frame and use the absolute source + # timeline. The request endpoint is in seconds; the training manifest used + # the same millisecond grid and ``precomputed_absolute_pts`` backend. + absolute_ms = np.rint((float(req.start_time) + indices / base_fps) * 1000.0).astype(np.int64) + metadata = { + "fps": 1000, + "frames_indices": absolute_ms.tolist(), + "total_num_frames": int(absolute_ms.max()) + 1, + "width": int(frames.shape[2]), + "height": int(frames.shape[1]), + "duration": (int(absolute_ms.max()) + 1) / 1000.0, + "video_backend": "precomputed_absolute_pts", + } + return frames, metadata + + +def answer_one(model, processor, req: Request) -> str: + """Run the merged Qwen3.5 checkpoint on one official request.""" + frames, metadata = your_preprocessing(clip_path_for(req), req) + messages = [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": [{"type": "video"}, {"type": "text", "text": req.question}]}, + ] + text = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True, enable_thinking=False + ) + inputs = processor( + text=[text], + videos=[frames], + video_metadata=[metadata], + do_sample_frames=False, + return_tensors="pt", + ) + model_device = next(model.parameters()).device + inputs = inputs.to(model_device) + with torch.inference_mode(): + generated = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=False) + new_tokens = generated[:, inputs["input_ids"].shape[1] :] + answer = processor.batch_decode( + new_tokens, skip_special_tokens=True, clean_up_tokenization_spaces=False + )[0].strip() + # The official evaluator consumes the content verbatim; only remove the + # optional reasoning wrapper if a checkpoint emits it despite disabled + # thinking, preserving all answer text otherwise. + if "" in answer: + answer = answer.split("", 1)[1].strip() + return answer + + +def run() -> int: + t_start = time.monotonic() + log.info("=== ORena SAVE FOCUS — inference start ===") + + # ── Device ──────────────────────────────────────────────────────────────── + # On the platform you get one GPU, chosen by you when you create the algorithm: + # an NVIDIA L40S (48 GB VRAM, 4 vCPUs, 32 GiB RAM) or an NVIDIA RTX PRO 6000 + # Blackwell (96 GB VRAM, 8 vCPUs, 64 GiB RAM). This image runs on both — see + # the hardware table in the README. + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + log.info("Device: %s", device) + log_environment(device) + + # ── Inputs ───────────────────────────────────────────────────────────────── + log.info("--- Loading inputs ---") + for name in ("batch.json", "request.json", "FO_definitions.json"): + p = INPUT_PATH / name + status = f"{p.stat().st_size} bytes" if p.exists() else "MISSING" + log.info(" %-24s %s", name, status) + log.info(" %-24s %s", f"{VIDEO_DIR.name}/", f"{len(list(VIDEO_DIR.glob('*.mp4')))} clip(s)" if VIDEO_DIR.is_dir() else "MISSING") + + # request.json holds the whole batch: answer every question it lists. + requests = load_requests(INPUT_PATH / "request.json") + if not requests: + log.error("request.json contains no requests") + return 1 + log.info("Batch of %d question(s)", len(requests)) + + # FO_definitions.json holds plain text, JSON-encoded; json.loads() gives the + # text back. The definitions are the same for every question. + fo_definitions = json.loads((INPUT_PATH / "FO_definitions.json").read_text()) + log.info("FO definitions: %d chars", len(fo_definitions)) + + # ── Model — loaded ONCE for the whole batch ──────────────────────────────── + # Keep this outside the per-question loop. Loading the model is paid once per + # run, and the latency budget grants a setup allowance to cover it. + log.info("--- Loading model (once for the batch) ---") + if not MODEL_PATH.exists(): + log.error("Merged model directory not found: %s", MODEL_PATH) + return 1 + log.info("Merged model path: %s", MODEL_PATH) + processor = AutoProcessor.from_pretrained(MODEL_PATH, local_files_only=True) + model = Qwen3_5ForConditionalGeneration.from_pretrained( + MODEL_PATH, + torch_dtype=torch.bfloat16, + device_map="auto", + low_cpu_mem_usage=True, + local_files_only=True, + ) + model.eval() + n_params = sum(p.numel() for p in model.parameters()) + log.info("Model loaded and set to eval mode: %d parameters (setup %.2f s)", n_params, time.monotonic() - t_start) + + if device.type == "cuda": + allocated = torch.cuda.memory_allocated(0) / 1024**2 + reserved = torch.cuda.memory_reserved(0) / 1024**2 + log.info(" VRAM after model load: %.1f MiB allocated / %.1f MiB reserved", allocated, reserved) + + # ── Inference — one pass per question ────────────────────────────────────── + log.info("--- Running inference over %d question(s) ---", len(requests)) + responses = [] + n_failed = 0 + t_batch = time.monotonic() + + for i, req in enumerate(requests, start=1): + log.info( + "[%d/%d] qID=%s videoID=%s window=[%.2f, %.2f] s (%.2f s)", + i, len(requests), req.qID, req.videoID, req.start_time, req.end_time, req.duration, + ) + t0 = time.monotonic() + try: + answer = answer_one(model, processor, req) + except Exception: + # Keep going: one failed question should not cost you the rest of the + # batch. An empty answer is simply scored incorrect. + n_failed += 1 + log.exception("[%d/%d] qID=%s failed; emitting empty answer", i, len(requests), req.qID) + answer = "" + latency = time.monotonic() - t0 + + # Informational only. The platform times the run itself and scores against + # that measurement — see "The latency budget" in the module docstring. + responses.append(Response(qID=req.qID, content=answer, latency=latency)) + log.info("[%d/%d] qID=%s answered in %.4f s: %r", i, len(requests), req.qID, latency, answer) + + batch_seconds = time.monotonic() - t_batch + log.info( + "Inference complete: %d answered (%d failed) in %.2f s (%.3f s/question)", + len(responses), n_failed, batch_seconds, batch_seconds / max(len(responses), 1), + ) + + if device.type == "cuda": + allocated = torch.cuda.memory_allocated(0) / 1024**2 + log.debug(" VRAM after inference: %.1f MiB allocated", allocated) + + # ── Output ───────────────────────────────────────────────────────────────── + log.info("--- Writing output ---") + OUTPUT_PATH.mkdir(parents=True, exist_ok=True) + output_path = OUTPUT_PATH / "answer.json" + save_items(responses, output_path) + log.info("Wrote %d response(s) to %s", len(responses), output_path) + log.info("=== inference done in %.2f s total ===", time.monotonic() - t_start) + return 0 + + +if __name__ == "__main__": + raise SystemExit(run()) diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/requirements.txt b/pipelines/Tr006_V2/docker/segment-algorithm/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..c8cf04d7f1ebe36599bc5eff6b7dd8c478cb4d11 --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/requirements.txt @@ -0,0 +1,21 @@ +# Pins the torch already in the base image — the recommended CUDA 12.8 build (see +# the Dockerfile for why it is the default). Nothing is downloaded: pip treats the +# preinstalled build as satisfying this. Its only job is to turn "a dependency of +# mine wants a different torch" into an immediate, readable pip error naming that +# dependency, rather than a silent swap to the CUDA 13.0 wheel that cannot run on +# the L40S. Changing torch is allowed — if you do, edit this pin knowingly, and the +# Dockerfile's guard will tell you whether the result still suits the GPU you target. +torch==2.11.0 + +orena-focus>=0.3.4 + +# Direct imports of inference.py. They also arrive as transitive dependencies +# of orena-focus, but anything imported directly is declared explicitly. +decord>=0.6.0 +numpy>=1.26.0 + +# Qwen3.5 support (including Qwen3_5ForConditionalGeneration and the +# timestamp-aware Qwen3VLProcessor) landed after the 4.x line. The merged +# checkpoint was produced and verified with this exact version. +transformers==5.15.0 +accelerate>=1.10.1 diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/resources/README.md b/pipelines/Tr006_V2/docker/segment-algorithm/resources/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cd8b7d09f4cf0687356197e0aaabbbabf6265b75 --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/resources/README.md @@ -0,0 +1,43 @@ +# `resources/` — your model lives here + +Everything in this folder is **copied into the container image** at build time and +is available at inference. This is where your trained model belongs. + +| File | What it is | Replace it? | +|:--|:--|:--| +| `model.py` | the model definition (`DummyModel`, a trivial CNN) | **yes** — with your architecture | +| `dummy_weights.pt` | a tiny checkpoint so the template runs out of the box | **yes** — with your trained weights | + +`inference.py` imports from here: + +```python +from resources.model import DummyModel + +WEIGHTS_PATH = RESOURCES_PATH / "dummy_weights.pt" +``` + +Keep that wiring, or point it at whatever you name your own files. + +## Things to get right + +**There is no internet at inference time.** The container runs with networking +disabled, so nothing can be downloaded at runtime — no model hubs, no checkpoints, +no tokenizer files, no fonts. Everything your model touches must be inside the +image. If you use a library that lazily fetches weights on first call (many +`transformers`, `timm`, and `open_clip` entry points do), pre-download during the +Docker build and load from a local path instead. + +**Load your model once.** One container run answers a whole batch of questions, so +the checkpoint is loaded once and reused across every question in that batch. Keep +model loading outside the per-question loop in `inference.py` — that is where the +setup allowance in the latency budget comes from. + +**Large checkpoints need splitting.** A single image layer cannot exceed **50 GB**, +and each Docker `COPY` is one layer. If your weights are bigger than that, split +them into chunks, copy them with several `COPY` instructions, and reassemble at +runtime — see the comment in the `Dockerfile`. + +**Big weights make big images.** The whole image is pulled before your code starts. +That pull is infrastructure overhead and is not charged against your latency +budget, but it does slow every submission cycle, so keep the image only as large as +it needs to be. diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/resources/dummy_weights.pt b/pipelines/Tr006_V2/docker/segment-algorithm/resources/dummy_weights.pt new file mode 100644 index 0000000000000000000000000000000000000000..23c98c69376a584b3e70d3b26e12be4c29dd2b81 --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/resources/dummy_weights.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:907b28b86b4d3262750ad946ab0384a40887e239b5be4027a53592817206dc2a +size 26973 diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/resources/model.py b/pipelines/Tr006_V2/docker/segment-algorithm/resources/model.py new file mode 100644 index 0000000000000000000000000000000000000000..82fcaf503a445d76ea86c6db23ba9115ba1cdc4d --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/resources/model.py @@ -0,0 +1,20 @@ +import torch.nn as nn + + +class DummyModel(nn.Module): + """Minimal CNN for verifying GPU execution. Replace with your real model.""" + + def __init__(self): + super().__init__() + self.features = nn.Sequential( + nn.Conv2d(3, 16, kernel_size=5, padding=1), + nn.ReLU(inplace=True), + nn.AdaptiveAvgPool2d(8), + nn.Conv2d(16, 32, kernel_size=3, padding=1), + nn.ReLU(inplace=True), + nn.AdaptiveAvgPool2d(1), + ) + self.classifier = nn.Linear(32, 2) + + def forward(self, x): + return self.classifier(self.features(x).flatten(1)) diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/resources/model_provenance.json b/pipelines/Tr006_V2/docker/segment-algorithm/resources/model_provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..7ed4c146f6307be9630ae9e89a8b61f50b888d4c --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/resources/model_provenance.json @@ -0,0 +1,17 @@ +{ + "model_family": "Qwen3.5-9B", + "weight_form": "merged_lora", + "lora_merged": true, + "source_base": "/cpfs01/projects-HDD/cfff-906dc71fafda_HDD/gbw_21307130160/cache/modelscope/models/Qwen--Qwen3.5-9B/snapshots/master", + "source_adapter_checkpoint": "/cpfs01/projects-HDD/cfff-906dc71fafda_HDD/gbw_21307130160/challenges/FARM-ORENA/experiments/Qwen3.5-9B/Tr/Tr006_offline_fixed80_absmeta_v2_lora64_bs4_10ep_20260828/results/checkpoint_archive/checkpoint-6013", + "source_step": 6013, + "source_epoch": 7.0, + "merge_report": "model/merge_report.json", + "merge_validation": { + "remaining_lora_modules_after_merge": 0, + "adapted_components": ["language_model", "visual", "visual_merger"] + }, + "runtime_model_class": "Qwen3_5ForConditionalGeneration", + "runtime_processor_class": "Qwen3VLProcessor", + "runtime_transformers": "5.15.0" +} diff --git a/pipelines/Tr006_V2/docker/segment-algorithm/test/.gitignore b/pipelines/Tr006_V2/docker/segment-algorithm/test/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..ea1472ec1f38b2d1410951194a02f8fb9f1c820a --- /dev/null +++ b/pipelines/Tr006_V2/docker/segment-algorithm/test/.gitignore @@ -0,0 +1 @@ +output/ diff --git a/pipelines/Tr006_V2/expected_outputs/v1/provenance.json b/pipelines/Tr006_V2/expected_outputs/v1/provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..7e48d7804f0e033d7e390bd383cf5897b0e79ecd --- /dev/null +++ b/pipelines/Tr006_V2/expected_outputs/v1/provenance.json @@ -0,0 +1,59 @@ +{ + "version": "fixed_frames_v1", + "created_at_utc": "2026-08-26T23:09:17.066162+00:00", + "command": "/cpfs01/projects-HDD/cfff-906dc71fafda_HDD/gbw_21307130160/challenges/FARM-ORENA/scripts/preprocess_segmentqa_fixed_frames.py --data-root /cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge --output-dir /cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/derived/orena_segmentqa_fixed80_420x224_absmeta_v1/clean --splits train test --frames 80 --width 420 --height 224 --workers 8 --resume", + "inputs": { + "heico-focus-vqa": { + "train": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/heico-focus-vqa/data/segment/train.parquet", + "test": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/heico-focus-vqa/data/segment/test.parquet", + "videos": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/heico-focus-vqa/videos" + }, + "lapchole-focus-vqa": { + "train": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/lapchole-focus-vqa/data/segment/train.parquet", + "test": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/lapchole-focus-vqa/data/segment/test.parquet", + "videos": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/lapchole-focus-vqa/videos" + } + }, + "splits": [ + "train", + "test" + ], + "question_count": 20000, + "unique_segment_count": 18478, + "failed_count": 0, + "frame_count": 80, + "sampling_rule": "uniform PTS bin centers in [start,end), nearest real source PTS", + "resize": "proportional letterbox 420x224", + "encoding": { + "container": "mp4", + "codec": "H.264/libx264", + "audio": "none", + "pix_fmt": "yuv420p" + }, + "metadata_contract": { + "fps": 1000, + "frames_indices": "round(abs_pts*1000)", + "total_num_frames": "max(frames_indices)+1", + "width": 420, + "height": 224, + "duration": "total_num_frames/fps", + "video_backend": "precomputed_absolute_pts" + }, + "system_prompt_source": { + "author_repo_commit": "3fc0850f9f76261c529e75c78ad2272558ff9e4e", + "source": "https://github.com/AzJ24/Orena/blob/3fc0850f9f76261c529e75c78ad2272558ff9e4e/orena_sft/prompts.py", + "function": "build_system_prompt", + "parameters": { + "include_definitions": false, + "style": "direct", + "track": "segment" + }, + "foreign_objects_source": "orena-focus 0.3.4 wheel: focus/foreign_objects.py" + }, + "system_prompt": "You are a surgical video analysis assistant. You are shown a short clip from a laparoscopic procedure and asked a single question about it.\n\nThe clip is given as a sequence of frames in chronological order. Each frame is preceded by its timestamp in the source video, written as -- a frame marked <1234.5 seconds> was taken 1234.5 seconds (00:20:34) after the video began. Consecutive frames may be several seconds apart, so an event can happen between two of them.\n\nA foreign object (FO) is any object fully introduced into the patient's body\ncavity during surgery that must be retrieved or accounted for. Importantly,\nstandard surgical instruments that remain connected to the external environment\n(e.g., graspers, scissors, trocars, staplers, cameras) are not considered foreign\nobjects. Furthermore, we exclude detachable parts of surgical instruments,\nparticularly anvil components of staplers.\n\nThe foreign object classes are exactly: Sponge, Clip, Specimen Bag, Silicone Loop, External Drain, Needle, Gallstone, Specimen, Mesh, Absorbable Hemostatic Agent.\n\nReply with the answer and nothing else -- no reasoning, no preamble, no\nexplanation, no restating the question. A single short line.\n\nRules for the answer:\n- Write the value only. No sentence, no explanation, no units, no trailing\n period, and never repeat the question.\n- Asks yes or no -> write exactly: yes or no\n- Asks how many / for a count -> write digits only, e.g. 0 or 1 or 2.\n- Asks which foreign object class(es) -> write class names exactly as spelled\n in the list above, comma-separated (e.g. Clip, Sponge), or exactly: none\n Never answer with a generic description such as \"surgical instrument\".\n- Asks when something happens -> write hh:mm:ss, read off the frame timestamps.\n Those timestamps count from the START OF THE VIDEO, not from the start of this\n clip, so <1234.5 seconds> is 00:20:34. Convert carefully: 3600 seconds is one\n hour, 60 seconds is one minute. If the event falls between two frames, give\n your best estimate between their timestamps rather than snapping to one.\n For several time points, separate them with commas: 00:09:19, 00:12:44\n- Asks for a percentage or a share -> write the number only, e.g. 40\n- Lists options to choose from -> copy exactly one of those options, verbatim.\n- Anything else -> a short phrase, at most a few words.\n\nIf you are unsure, still commit to your single best answer in the required\nform. An empty, hedged, or explanatory answer is scored as wrong.\n", + "output_dir": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/derived/orena_segmentqa_fixed80_420x224_absmeta_v1/clean", + "resume": true, + "workers": 8, + "limit": null, + "seed": 42 +} diff --git a/pipelines/Tr006_V2/expected_outputs/v1/validation_report.json b/pipelines/Tr006_V2/expected_outputs/v1/validation_report.json new file mode 100644 index 0000000000000000000000000000000000000000..288211efed92b24f89e338d03c49470f7ab6d918 --- /dev/null +++ b/pipelines/Tr006_V2/expected_outputs/v1/validation_report.json @@ -0,0 +1,11 @@ +{ + "status": "pass", + "expected_frame_count": 80, + "expected_width": 420, + "expected_height": 224, + "segments": 18478, + "success": 18478, + "missing": 0, + "failures": 0, + "checked_at_utc": "2026-08-26T23:09:17.066162+00:00" +} diff --git a/pipelines/Tr006_V2/expected_outputs/v2/provenance.json b/pipelines/Tr006_V2/expected_outputs/v2/provenance.json new file mode 100644 index 0000000000000000000000000000000000000000..93b7eee89cf3d0548cb64611f70bb690d036e076 --- /dev/null +++ b/pipelines/Tr006_V2/expected_outputs/v2/provenance.json @@ -0,0 +1,27 @@ +{ + "version": "fixed_frames_blueclean_v2", + "created_at_utc": "2026-08-27T03:55:02.119997+00:00", + "command": "scripts/preprocess_segmentqa_blueclean_v2.py --data-root /cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge --output-dir /cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/derived/orena_segmentqa_fixed80_420x224_absmeta_v2/clean --splits train test --frames 80 --width 420 --height 224 --workers 8 --min-train-valid-frames 16 --resume", + "data_root": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge", + "v1_dir": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/derived/orena_segmentqa_fixed80_420x224_absmeta_v1/clean", + "blue_intervals": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/heico-focus-vqa/blue_screen_temporal_annotations/data/intervals.csv", + "output_dir": "/cpfs01/projects-SSD/cfff-906dc71fafda_SSD/gbw_21307130160/dataset/challenge/derived/orena_segmentqa_fixed80_420x224_absmeta_v2/clean", + "splits": [ + "train", + "test" + ], + "frame_count": 80, + "resize": "proportional letterbox 420x224", + "encoding": { + "container": "mp4", + "codec": "H.264/libx264", + "audio": "none" + }, + "sampling_rule": "80 uniform bin centers over the non-blue complement inside the original [start,end)", + "min_train_valid_frames": 16, + "unaffected_copy": "cp --reflink=auto; byte-identical verification; no hardlinks", + "test_answers_used_for_decisions": false, + "resume": true, + "workers": 8, + "limit": null +} diff --git a/pipelines/Tr006_V2/expected_outputs/v2/validation_report.json b/pipelines/Tr006_V2/expected_outputs/v2/validation_report.json new file mode 100644 index 0000000000000000000000000000000000000000..044700872e6a6583457e9b024d7da6b654ef2e35 --- /dev/null +++ b/pipelines/Tr006_V2/expected_outputs/v2/validation_report.json @@ -0,0 +1,24 @@ +{ + "status": "pass", + "expected_frame_count": 80, + "expected_width": 420, + "expected_height": 224, + "segments_total": 18478, + "affected_temporal_overlap": 1215, + "unaffected": 17263, + "included_segments": 18478, + "excluded_train_segments": 0, + "excluded_train_questions": 0, + "train_questions": 13746, + "test_questions": 6254, + "test_unique_segments": 5661, + "actions": { + "copied": 17263, + "resampled_excluding_blue": 1214, + "retained_test_insufficient_valid": 1 + }, + "missing": 0, + "failures": 0, + "full_dataset_contract": true, + "checked_at_utc": "2026-08-27T03:55:02.119997+00:00" +} diff --git a/pipelines/Tr006_V2/manifest/pipeline_manifest.json b/pipelines/Tr006_V2/manifest/pipeline_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..994412fcf0f649a566f9b1aafee97de283b70ad4 --- /dev/null +++ b/pipelines/Tr006_V2/manifest/pipeline_manifest.json @@ -0,0 +1,60 @@ +{ + "bundle_id": "Tr006_V2_checkpoint-6013", + "track": "segment", + "project_commit": "b66edaeba3c484c431d785f95a7fc81cffa530c0", + "code_state": "dirty; exact files included in this bundle are authoritative", + "base_model": "Qwen3.5-9B snapshots/master; not included", + "adapter_path": "weights/Tr006_V2/checkpoint-6013", + "raw_data": { + "included": false, + "required_layout": "README.md#原始数据目录契约", + "parquet_sha256": { + "heico-focus-vqa/data/segment/train.parquet": "ee6b1385a708ebe284cca5368e918bc62d72fb3bad2d9fc342d7a4fe4562d4fe", + "heico-focus-vqa/data/segment/test.parquet": "da7891b52269d27b89e21237efdaea7f51b3ee3ad0f82573eebdf6fbbf6985e9", + "lapchole-focus-vqa/data/segment/train.parquet": "f2198a516e9bb051a9071c2a80c1389fc9f283d81357d4e52f4d280f18f6eb70", + "lapchole-focus-vqa/data/segment/test.parquet": "f4522f845ceb6a3847f45d90f0bba84aa7df98abc4dd5e4cfd4fb914ceccac95" + } + }, + "blue_intervals": { + "count": 323, + "sha256": "d523a529226f9bb8ec06147874a6321e1f5c201700d55a09fd8baecdb1240f03" + }, + "v1_expected": { + "segments": 18478, + "frames": 80, + "width": 420, + "height": 224, + "train_jsonl_sha256": "a84cb85db66665a7b37d7827104ae053c1ef710a748b56d5d90198a71078f7c2", + "segment_manifest_sha256": "1a61c9e60647af24f7019e9fa0209aed10d63e2a0bf9d7ccb5cccd75eb3d06a6" + }, + "v2_expected": { + "segments": 18478, + "affected_temporal_overlap": 1215, + "train_questions": 13746, + "test_questions": 6254, + "frames": 80, + "width": 420, + "height": 224, + "metadata_backend": "precomputed_absolute_pts", + "train_jsonl_sha256": "197800db3323b311a8ffb2c6459919d3b4db3ca1487b2f8790c7c5eacf16a4e2", + "segment_manifest_sha256": "6098c0abb1843be6ac24498914fb9c43c19d8bc4c2d034fbffb0d66456ad6414" + }, + "training": { + "transformers": "5.15.0", + "ms_swift": "4.5.0.dev0", + "epochs": 10, + "selected_checkpoint_step": 6013, + "selected_checkpoint_epoch": 7.0, + "frames": 80, + "resolution": "420x224 letterbox", + "precision": "bfloat16", + "learning_rate": "1e-4", + "lora_rank": 64, + "lora_alpha": 128 + }, + "submission_image_snapshot": { + "image_id": "sha256:3ea66ddf64a40c9d2389f7c78326ab3728d54d6bb287b0cfa27cdee391827c5f", + "runtime_backend": "transformers Qwen3_5ForConditionalGeneration", + "actual_gpu_interface_smoke_before_submission": false + } +} diff --git a/pipelines/Tr006_V2/manifest/pipeline_requirements.txt b/pipelines/Tr006_V2/manifest/pipeline_requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..9993e3ab0ab0781d8e57c9dd9803fe60e4790d02 --- /dev/null +++ b/pipelines/Tr006_V2/manifest/pipeline_requirements.txt @@ -0,0 +1,15 @@ +# Raw blue-screen detection +opencv-python==4.11.0.86 +numpy==2.4.6 + +# Segment fixed80 V1/V2 generation +av==18.1.0 +pandas==2.2.3 +Pillow==11.3.0 +pyarrow + +# Training/merge versions recorded for Tr006 +ms-swift==4.5.0.dev0 +transformers==5.15.0 +peft==0.19.1 +safetensors==0.8.0 diff --git a/pipelines/Tr006_V2/project/data_preprocessing/README.md b/pipelines/Tr006_V2/project/data_preprocessing/README.md new file mode 100644 index 0000000000000000000000000000000000000000..5fa37fcc17ce6799a5dbf99d0165b307ec4a351a --- /dev/null +++ b/pipelines/Tr006_V2/project/data_preprocessing/README.md @@ -0,0 +1,92 @@ +# 蓝屏时序预处理脚本说明 + +本目录用于检测胃肠镜视频中的冻结蓝屏区段、细化时间边界，并将结果整理为可发布的数据集。所有检测脚本均以原始视频为只读输入，不会裁剪、覆盖或移动源视频。 + +## 推荐流程 + +```text +视频目录 + -> process_blue_screen_dataset.py + -> results/{intervals.csv,videos.jsonl,per_video/,evidence/,summary.json} + -> 结果验收及 validation.json + -> prepare_blue_screen_hf_release.py + -> 可上传的 HF Dataset 目录 +``` + +正式批量处理优先使用 `process_blue_screen_dataset.py`。`blue_screen_detector.py`、`fast_blue_screen_detector.py` 和 `refine_blue_intervals.py` 主要用于方法验证、单视频诊断和边界复核，不应替代正式批处理入口。 + +## 脚本清单 + +### `process_blue_screen_dataset.py` + +正式的全数据集粗到细检测入口。 + +- 输入：包含 `.avi`、`.mp4`、`.mov` 或 `.mkv` 视频的目录。 +- 粗检：顺序解码视频，每秒抽取一次圆形内镜 ROI，对 ROI 像素计算哈希和蓝色比例，定位连续静止蓝屏候选区段。 +- 精修：只对候选区段及其前后缓冲区逐帧扫描，以原生 FPS 确定准确的开始帧和结束帧。 +- 判定：ROI 解码帧哈希完全相同，同时蓝色比例达到阈值并满足最短持续时间。 +- 支持 `--resume`，可复用已经完成的逐视频 JSON 结果。 +- 输出： + - `intervals.csv`：每个蓝屏区段一行，是主要时序标注。 + - `videos.jsonl`：每个视频一行，包含视频信息和全部区段。 + - `per_video/.json`：逐视频检查点和详细结果。 + - `evidence//`：每个区段的前一帧、起始帧、结束帧和后一帧证据图。 + - `summary.json`：全数据集汇总统计。 + +正式 HeiCo 运行使用的主要参数为：1 秒粗采样、最短静止 6 秒、蓝色比例阈值 0.45、候选区段前后各扩展 5 秒、粗检尺寸 96、精修尺寸 160。 + +### `fast_blue_screen_detector.py` + +单视频快速哈希检测工具，用于验证“蓝屏期间抽样帧是否完全相同”的假设。 + +- 顺序解码视频，但默认只每秒生成一个 ROI 哈希和蓝色比例。 +- 将连续相同哈希且持续时间足够长的样本合并为粗粒度区段。 +- 可通过重复传入 `--reference START END` 与用户提供的参考区段计算近似秒级 precision/recall。 +- 输出单个 JSON；时间边界精度由 `--sample-seconds` 决定，不能代替逐帧精修结果。 + +### `refine_blue_intervals.py` + +针对已知粗区段的单视频逐帧边界精修工具。 + +- 通过一个或多个 `--interval START END` 接收粗区段。 +- 在粗区段前后增加 `--pad` 秒扫描范围。 +- 逐帧计算圆形 ROI 蓝色比例，并通过连续确认帧抑制孤立噪声。 +- 输出精修后的帧号、秒级时间、持续时长和平均蓝色比例。 +- 适合人工复核或早期实验；正式批处理采用 `process_blue_screen_dataset.py` 内置的“完全重复帧哈希 + 蓝色比例”精修实现。 + +### `blue_screen_detector.py` + +早期单视频候选检测工具，用于验证基于颜色和纹理的蓝屏判定方法。 + +- 自动估计圆形内镜 ROI，忽略外围黑色区域。 +- 稀疏抽帧后计算蓝色比例、拉普拉斯纹理强度、平均色相和有效像素比例。 +- 将高蓝色比例、低纹理且持续时间足够长的样本合并为候选区段。 +- 输出详细 JSON 和区段 CSV。 +- 该方法没有完全相同帧哈希约束，结果属于候选范围，不作为当前正式蓝屏时序标注。 + +### `prepare_blue_screen_hf_release.py` + +将已经验收通过的运行结果整理为可移植的 Hugging Face Dataset 发布目录。 + +- 要求源运行包含 `run.json`、`results/intervals.csv`、`videos.jsonl`、`summary.json`、`validation.json` 和 `evidence/`。 +- 只有 `validation.json` 状态为 `PASS` 时才允许打包。 +- 删除结果中的服务器绝对视频路径，改用 `video_id` 和 `video_filename`。 +- 复制边界证据图并生成便携的 provenance、summary、validation 和 release manifest。 +- 对视频数、区段数和证据图数量进行一致性核验。 +- 为避免误覆盖，目标输出目录已存在时会直接停止。 + +## 时间区间语义 + +蓝屏区段统一使用半开区间 `[start_seconds, end_seconds)`：开始时间包含在区段内，结束时间指向蓝屏结束后的第一个时间位置。逐帧结果同时提供 `start_frame` 和 `end_frame`，其中帧号边界为首尾均包含。 + +## 依赖与运行环境 + +- Python 3 +- OpenCV (`cv2`) +- NumPy + +FVQA 项目中的正式运行应使用 `fvqa-harness.yaml` 指定的 Conda 环境。蓝屏检测只使用 CPU，不需要占用 GPU。 + +## 当前正式结果 + +HeiCo 30 个视频的已验收结果位于 HF Dataset 仓库 `Gbone3176/orena-seg` 的 `blue_screen_temporal_annotations/` 子目录：共 323 个蓝屏区段和 1,291 张边界证据图。该结果通过了文件覆盖、逐帧解码、区段边界和汇总一致性检查；这里的 `PASS` 表示产物一致性验收通过，不代表基于独立人工金标准计算的检测准确率。 diff --git a/pipelines/Tr006_V2/project/data_preprocessing/process_blue_screen_dataset.py b/pipelines/Tr006_V2/project/data_preprocessing/process_blue_screen_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..8666cf4ee5e3f311c5f3223055048369c97da391 --- /dev/null +++ b/pipelines/Tr006_V2/project/data_preprocessing/process_blue_screen_dataset.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +"""Coarse-to-fine frozen blue-screen detection for endoscopy videos. + +Stage 1 decodes each video sequentially and hashes one circular-ROI frame per +second. Stage 2 re-reads only candidate windows and locates exact repeated-frame +runs at native FPS. Inputs are read-only; all results go below --output-dir. +""" +from __future__ import annotations + +import argparse +import csv +import hashlib +import json +import socket +import time +from datetime import datetime, timezone +from pathlib import Path + +import cv2 +import numpy as np + +VIDEO_SUFFIXES = {".avi", ".mp4", ".mov", ".mkv"} + + +def utc_now() -> str: + return datetime.now(timezone.utc).astimezone().isoformat() + + +def hhmmss(seconds: float) -> str: + ms = int(round(seconds * 1000)) + hours, rem = divmod(ms, 3_600_000) + minutes, rem = divmod(rem, 60_000) + secs, millis = divmod(rem, 1000) + return f"{hours:02d}:{minutes:02d}:{secs:02d}.{millis:03d}" + + +def roi_mask(height: int, width: int) -> np.ndarray: + yy, xx = np.ogrid[:height, :width] + return (((xx - width / 2) / (width * 0.47)) ** 2 + + ((yy - height / 2) / (height * 0.47)) ** 2 <= 1) + + +def frame_features(frame: np.ndarray, mask: np.ndarray, size: int) -> tuple[str, float]: + small = cv2.resize(frame, (size, size), interpolation=cv2.INTER_AREA) + small_mask = cv2.resize(mask.astype(np.uint8), (size, size), + interpolation=cv2.INTER_NEAREST).astype(bool) + pixels = small[small_mask] + digest = hashlib.md5(pixels.tobytes()).hexdigest() + hsv = cv2.cvtColor(small, cv2.COLOR_BGR2HSV) + h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2] + valid = small_mask & (v > 18) + if not valid.any(): + return digest, 0.0 + blue = valid & (h >= 88) & (h <= 140) & (s >= 55) & (v >= 35) + return digest, float(blue.sum() / valid.sum()) + + +def coarse_scan(video: Path, sample_seconds: float, min_static_seconds: float, + blue_threshold: float, size: int) -> tuple[dict, list[dict]]: + cap = cv2.VideoCapture(str(video)) + if not cap.isOpened(): + raise RuntimeError("cannot open video") + fps = float(cap.get(cv2.CAP_PROP_FPS) or 0) + frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) or 0) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) or 0) + if fps <= 0 or frame_count <= 0 or width <= 0 or height <= 0: + cap.release() + raise RuntimeError(f"invalid metadata fps={fps} frames={frame_count} size={width}x{height}") + mask = roi_mask(height, width) + step = max(1, int(round(fps * sample_seconds))) + samples: list[dict] = [] + frame_idx = 0 + next_sample = 0 + while True: + ok, frame = cap.read() + if not ok: + break + if frame_idx >= next_sample: + digest, ratio = frame_features(frame, mask, size) + samples.append({"frame": frame_idx, "time": frame_idx / fps, + "hash": digest, "blue_ratio": ratio}) + next_sample += step + frame_idx += 1 + cap.release() + if frame_idx < frame_count * 0.99: + raise RuntimeError(f"early decode stop decoded={frame_idx} metadata_frames={frame_count}") + + min_samples = max(2, int(np.ceil(min_static_seconds / sample_seconds))) + candidates: list[dict] = [] + run_start = 0 + for i in range(1, len(samples) + 1): + if i == len(samples) or samples[i]["hash"] != samples[run_start]["hash"]: + chunk = samples[run_start:i] + mean_blue = float(np.mean([x["blue_ratio"] for x in chunk])) + if len(chunk) >= min_samples and mean_blue >= blue_threshold: + candidates.append({ + "start_seconds": chunk[0]["time"], + "end_seconds": chunk[-1]["time"] + sample_seconds, + "sample_count": len(chunk), + "mean_blue_ratio": mean_blue, + "hash": chunk[0]["hash"], + }) + run_start = i + metadata = {"fps": fps, "frame_count": frame_count, "width": width, + "height": height, "duration_seconds": frame_count / fps, + "decoded_frames": frame_idx, "coarse_samples": len(samples)} + return metadata, candidates + + +def refine_candidate(video: Path, metadata: dict, candidate: dict, pad_seconds: float, + min_static_seconds: float, blue_threshold: float, + size: int) -> list[dict]: + fps = metadata["fps"] + first = max(0, int(np.floor((candidate["start_seconds"] - pad_seconds) * fps))) + last = min(metadata["frame_count"] - 1, + int(np.ceil((candidate["end_seconds"] + pad_seconds) * fps))) + cap = cv2.VideoCapture(str(video)) + cap.set(cv2.CAP_PROP_POS_FRAMES, first) + mask = roi_mask(metadata["height"], metadata["width"]) + samples: list[dict] = [] + frame_idx = first + while frame_idx <= last: + ok, frame = cap.read() + if not ok: + break + digest, ratio = frame_features(frame, mask, size) + samples.append({"frame": frame_idx, "hash": digest, "blue_ratio": ratio}) + frame_idx += 1 + cap.release() + + min_frames = max(2, int(np.ceil(min_static_seconds * fps))) + runs: list[dict] = [] + run_start = 0 + for i in range(1, len(samples) + 1): + if i == len(samples) or samples[i]["hash"] != samples[run_start]["hash"]: + chunk = samples[run_start:i] + mean_blue = float(np.mean([x["blue_ratio"] for x in chunk])) + if len(chunk) >= min_frames and mean_blue >= blue_threshold: + start_frame = chunk[0]["frame"] + end_frame = chunk[-1]["frame"] + runs.append({ + "start_frame": start_frame, + "end_frame": end_frame, + "start_seconds": start_frame / fps, + "end_seconds": (end_frame + 1) / fps, + "duration_seconds": len(chunk) / fps, + "start_time": hhmmss(start_frame / fps), + "end_time": hhmmss((end_frame + 1) / fps), + "frame_count": len(chunk), + "mean_blue_ratio": mean_blue, + "exact_repeated_hash": chunk[0]["hash"], + "confidence": "high", + }) + run_start = i + return runs + + +def merge_duplicate_runs(runs: list[dict]) -> list[dict]: + by_key: dict[tuple[int, int], dict] = {} + for run in runs: + by_key[(run["start_frame"], run["end_frame"])] = run + return sorted(by_key.values(), key=lambda x: x["start_frame"]) + + +def save_evidence(video: Path, intervals: list[dict], output_dir: Path, fps: float) -> None: + if not intervals: + return + output_dir.mkdir(parents=True, exist_ok=True) + cap = cv2.VideoCapture(str(video)) + for index, interval in enumerate(intervals, 1): + points = { + "before": max(0, interval["start_frame"] - 1), + "start": interval["start_frame"], + "end": interval["end_frame"], + "after": interval["end_frame"] + 1, + } + evidence = {} + for label, frame_idx in points.items(): + cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) + ok, frame = cap.read() + if ok: + name = f"interval_{index:03d}_{label}_f{frame_idx}.jpg" + cv2.imwrite(str(output_dir / name), frame) + evidence[label] = str((output_dir / name).name) + interval["evidence_frames"] = evidence + cap.release() + + +def write_aggregate(output_dir: Path, records: list[dict], started_at: str, + elapsed: float, input_dir: Path, config: dict) -> None: + intervals = [] + for record in records: + for idx, interval in enumerate(record.get("intervals", []), 1): + intervals.append({"video_id": record["video_id"], + "video_path": record["video_path"], + "interval_index": idx, **interval}) + fields = ["video_id", "video_path", "interval_index", "start_frame", "end_frame", + "start_seconds", "end_seconds", "duration_seconds", "start_time", "end_time", + "frame_count", "mean_blue_ratio", "confidence", "exact_repeated_hash"] + with (output_dir / "intervals.csv").open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fields, extrasaction="ignore") + writer.writeheader(); writer.writerows(intervals) + with (output_dir / "videos.jsonl").open("w", encoding="utf-8") as handle: + for record in records: + handle.write(json.dumps(record, ensure_ascii=False) + "\n") + summary = { + "status": "completed" if all(r["status"] == "completed" for r in records) else "failed", + "started_at": started_at, "finished_at": utc_now(), "elapsed_seconds": elapsed, + "host": socket.gethostname(), "input_dir": str(input_dir), "config": config, + "video_count": len(records), + "completed_video_count": sum(r["status"] == "completed" for r in records), + "failed_video_count": sum(r["status"] == "failed" for r in records), + "videos_with_blue_screen": sum(bool(r.get("intervals")) for r in records), + "interval_count": len(intervals), + "total_video_seconds": sum(r.get("duration_seconds", 0) for r in records), + "total_blue_screen_seconds": sum(x["duration_seconds"] for x in intervals), + } + (output_dir / "summary.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--input-dir", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--sample-seconds", type=float, default=1.0) + parser.add_argument("--min-static-seconds", type=float, default=6.0) + parser.add_argument("--blue-ratio", type=float, default=0.45) + parser.add_argument("--pad-seconds", type=float, default=5.0) + parser.add_argument("--coarse-size", type=int, default=96) + parser.add_argument("--fine-size", type=int, default=160) + parser.add_argument("--resume", action="store_true") + args = parser.parse_args() + videos = sorted(p for p in args.input_dir.rglob("*") if p.suffix.lower() in VIDEO_SUFFIXES) + if not videos: + raise SystemExit("no videos found") + args.output_dir.mkdir(parents=True, exist_ok=True) + per_video_dir = args.output_dir / "per_video" + evidence_root = args.output_dir / "evidence" + per_video_dir.mkdir(exist_ok=True) + started_at = utc_now(); t0 = time.perf_counter() + config = {"sample_seconds": args.sample_seconds, + "min_static_seconds": args.min_static_seconds, + "blue_ratio": args.blue_ratio, "pad_seconds": args.pad_seconds, + "coarse_size": args.coarse_size, "fine_size": args.fine_size, + "definition": "exact repeated decoded ROI frames with blue ratio above threshold"} + records = [] + for number, video in enumerate(videos, 1): + video_id = video.stem.split(" ", 1)[0] + result_path = per_video_dir / f"{video_id}.json" + if args.resume and result_path.exists(): + records.append(json.loads(result_path.read_text(encoding="utf-8"))) + print(f"[{number}/{len(videos)}] resume {video.name}", flush=True) + continue + print(f"[{number}/{len(videos)}] start {video.name} {utc_now()}", flush=True) + item_t0 = time.perf_counter() + record = {"video_id": video_id, "video_path": str(video), "status": "failed", + "started_at": utc_now(), "intervals": []} + try: + metadata, candidates = coarse_scan(video, args.sample_seconds, + args.min_static_seconds, + args.blue_ratio, args.coarse_size) + fine_runs = [] + for candidate in candidates: + fine_runs.extend(refine_candidate(video, metadata, candidate, + args.pad_seconds, + args.min_static_seconds, + args.blue_ratio, args.fine_size)) + fine_runs = merge_duplicate_runs(fine_runs) + save_evidence(video, fine_runs, evidence_root / video_id, metadata["fps"]) + blue_seconds = sum(x["duration_seconds"] for x in fine_runs) + record.update(metadata) + record.update({"status": "completed", "coarse_candidate_count": len(candidates), + "interval_count": len(fine_runs), "blue_screen": bool(fine_runs), + "blue_screen_seconds": blue_seconds, + "blue_screen_ratio": blue_seconds / metadata["duration_seconds"], + "intervals": fine_runs}) + except Exception as exc: + record["error"] = f"{type(exc).__name__}: {exc}" + record["finished_at"] = utc_now() + record["elapsed_seconds"] = time.perf_counter() - item_t0 + result_path.write_text(json.dumps(record, ensure_ascii=False, indent=2), encoding="utf-8") + records.append(record) + write_aggregate(args.output_dir, records, started_at, + time.perf_counter() - t0, args.input_dir, config) + print(f"[{number}/{len(videos)}] {record['status']} intervals={len(record['intervals'])} elapsed={record['elapsed_seconds']:.1f}s", flush=True) + write_aggregate(args.output_dir, records, started_at, + time.perf_counter() - t0, args.input_dir, config) + if any(r["status"] != "completed" for r in records): + raise SystemExit(2) + + +if __name__ == "__main__": + main() diff --git a/pipelines/Tr006_V2/project/ms-swift/LICENSE b/pipelines/Tr006_V2/project/ms-swift/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..261eeb9e9f8b2b4b0d119366dda99c6fd7d35c64 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/pipelines/Tr006_V2/project/ms-swift/MANIFEST.in b/pipelines/Tr006_V2/project/ms-swift/MANIFEST.in new file mode 100644 index 0000000000000000000000000000000000000000..4f918468371e898e949aeeba02cfe0e5203b0785 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/MANIFEST.in @@ -0,0 +1 @@ +recursive-include requirements *.txt diff --git a/pipelines/Tr006_V2/project/ms-swift/requirements.txt b/pipelines/Tr006_V2/project/ms-swift/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..0832e6ab4d42f6181fbfb354c0d88b1809980afe --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/requirements.txt @@ -0,0 +1 @@ +-r requirements/framework.txt diff --git a/pipelines/Tr006_V2/project/ms-swift/requirements/framework.txt b/pipelines/Tr006_V2/project/ms-swift/requirements/framework.txt new file mode 100644 index 0000000000000000000000000000000000000000..6829b7692b54445e63baf6a8b3cf9c077d73dff1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/requirements/framework.txt @@ -0,0 +1,39 @@ +accelerate +addict +aiohttp +attrdict +binpacking +charset_normalizer +cpm_kernels +dacite +datasets>=3.0,<4.8.5 +einops +fastapi +gradio>=3.40.0,<6.0 +importlib_metadata +json_repair +matplotlib +modelscope>=1.23 +nltk +numpy +openai +oss2 +pandas +peft>=0.11,<0.20 +pillow +PyYAML>=5.4 +requests +rouge +safetensors +scipy +sentencepiece +simplejson>=3.3.0 +sortedcontainers>=1.5.9 +tensorboard +tiktoken +tqdm +transformers>=4.33,<5.17.0 +transformers_stream_generator +trl>=0.15,<1.0 +uvicorn +zstandard diff --git a/pipelines/Tr006_V2/project/ms-swift/requirements/swanlab.txt b/pipelines/Tr006_V2/project/ms-swift/requirements/swanlab.txt new file mode 100644 index 0000000000000000000000000000000000000000..85381455c6b3bb922c0c4c23a82810e3ddb41f04 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/requirements/swanlab.txt @@ -0,0 +1 @@ +swanlab diff --git a/pipelines/Tr006_V2/project/ms-swift/setup.cfg b/pipelines/Tr006_V2/project/ms-swift/setup.cfg new file mode 100644 index 0000000000000000000000000000000000000000..0f8b9b019f2d63b728cf08047eee8674e4f8d5da --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/setup.cfg @@ -0,0 +1,32 @@ +[isort] +line_length = 120 +multi_line_output = 0 +known_standard_library = setuptools +known_first_party = swift +known_third_party = json,yaml +no_lines_before = STDLIB,LOCALFOLDER +default_section = THIRDPARTY + +[yapf] +BASED_ON_STYLE = pep8 +COLUMN_LIMIT = 120 +BLANK_LINE_BEFORE_NESTED_CLASS_OR_DEF = true +SPLIT_BEFORE_EXPRESSION_AFTER_OPENING_PAREN = true +SPLIT_BEFORE_ARITHMETIC_OPERATOR = true + +[codespell] +skip = *.ipynb +quiet-level = 3 +ignore-words-list = patten,nd,ty,mot,hist,formating,winn,gool,datas,wan,confids + +[flake8] +max-line-length = 120 +select = B,C,E,F,P,T4,W,B9 +ignore = F401,F403,F405,F821,W503,E251,W504,E126 +exclude = docs/src,*.pyi,.git,peft.py + +[darglint] +ignore=DAR101 + +[easy_install] +index-url=https://pypi.tuna.tsinghua.edu.cn/simple diff --git a/pipelines/Tr006_V2/project/ms-swift/setup.py b/pipelines/Tr006_V2/project/ms-swift/setup.py new file mode 100644 index 0000000000000000000000000000000000000000..6d74e4d0c5d9ebeadb4f711735f0af9e3e7e0c79 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/setup.py @@ -0,0 +1,166 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# !/usr/bin/env python +from setuptools import find_packages, setup + +import os +import re +from typing import List + + +def readme(): + with open('README.md', encoding='utf-8') as f: + content = f.read() + return content + + +version_file = 'swift/version.py' + + +def get_version(): + with open(version_file, 'r', encoding='utf-8') as f: + return re.search(r'^__version__\s*=\s*["\'](.+?)["\']', f.read(), re.M).group(1) + + +def parse_requirements(fname='requirements.txt', with_version=True): + """ + Parse the package dependencies listed in a requirements file but strips + specific versioning information. + + Args: + fname (str): path to requirements file + with_version (bool, default=False): if True include version specs + + Returns: + List[str]: list of requirements items + + CommandLine: + python -c "import setup; print(setup.parse_requirements())" + """ + import re + import sys + from os.path import exists + require_fpath = fname + + def parse_line(line): + """ + Parse information from a line in a requirements text file + """ + if line.startswith('-r '): + # Allow specifying requirements in other files + target = line.split(' ')[1] + relative_base = os.path.dirname(fname) + absolute_target = os.path.join(relative_base, target) + for info in parse_require_file(absolute_target): + yield info + else: + info = {'line': line} + if line.startswith('-e '): + info['package'] = line.split('#egg=')[1] + else: + # Remove versioning from the package + pat = '(' + '|'.join(['>=', '==', '>']) + ')' + parts = re.split(pat, line, maxsplit=1) + parts = [p.strip() for p in parts] + + info['package'] = parts[0] + if len(parts) > 1: + op, rest = parts[1:] + if ';' in rest: + # Handle platform specific dependencies + # http://setuptools.readthedocs.io/en/latest/setuptools.html#declaring-platform-specific-dependencies + version, platform_deps = map(str.strip, rest.split(';')) + info['platform_deps'] = platform_deps + else: + version = rest # NOQA + info['version'] = (op, version) + yield info + + def parse_require_file(fpath): + with open(fpath, 'r', encoding='utf-8') as f: + for line in f.readlines(): + line = line.strip() + if line.startswith('http'): + print('skip http requirements %s' % line) + continue + if line and not line.startswith('#') and not line.startswith('--'): + for info in parse_line(line): + yield info + elif line and line.startswith('--find-links'): + eles = line.split() + for e in eles: + e = e.strip() + if 'http' in e: + info = dict(dependency_links=e) + yield info + + def gen_packages_items(): + items = [] + deps_link = [] + if exists(require_fpath): + for info in parse_require_file(require_fpath): + if 'dependency_links' not in info: + parts = [info['package']] + if with_version and 'version' in info: + parts.extend(info['version']) + if not sys.version.startswith('3.4'): + # apparently package_deps are broken in 3.4 + platform_deps = info.get('platform_deps') + if platform_deps is not None: + parts.append(';' + platform_deps) + item = ''.join(parts) + items.append(item) + else: + deps_link.append(info['dependency_links']) + return items, deps_link + + return gen_packages_items() + + +if __name__ == '__main__': + install_requires, deps_link = parse_requirements('requirements.txt') + extra_requires = {} + all_requires = [] + extra_requires['megatron'], _ = parse_requirements('requirements/megatron.txt') + extra_requires['eval'], _ = parse_requirements('requirements/eval.txt') + extra_requires['swanlab'], _ = parse_requirements('requirements/swanlab.txt') + extra_requires['ray'], _ = parse_requirements('requirements/ray.txt') + all_requires.extend(install_requires) + all_requires.extend(extra_requires['eval']) + all_requires.extend(extra_requires['swanlab']) + all_requires.extend(extra_requires['ray']) + extra_requires['all'] = all_requires + + setup( + name='ms_swift', + version=get_version(), + description='Swift: Scalable lightWeight Infrastructure for Fine-Tuning', + long_description=readme(), + long_description_content_type='text/markdown', + author='DAMO ModelScope teams', + author_email='contact@modelscope.cn', + keywords=['transformers', 'LLM', 'lora', 'megatron', 'grpo', 'sft'], + url='https://github.com/modelscope/ms-swift', + packages=find_packages(exclude=('tests', 'tests.*')), + include_package_data=True, + package_data={'': ['utils/*', 'dataset/data/*.*', 'config/*.json', 'loss_scale/config/*.json']}, + python_requires='>=3.8.0', + classifiers=[ + 'Development Status :: 4 - Beta', + 'License :: OSI Approved :: Apache Software License', + 'Operating System :: OS Independent', + 'Programming Language :: Python :: 3', + 'Programming Language :: Python :: 3.8', + 'Programming Language :: Python :: 3.9', + 'Programming Language :: Python :: 3.10', + 'Programming Language :: Python :: 3.11', + 'Programming Language :: Python :: 3.12', + ], + license='Apache License 2.0', + tests_require=parse_requirements('requirements/tests.txt'), + install_requires=install_requires, + extras_require=extra_requires, + entry_points={ + 'console_scripts': ['swift=swift.cli.main:cli_main', 'megatron=swift.cli._megatron.main:cli_main'] + }, + dependency_links=deps_link, + zip_safe=False) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dbbc3e6f45efa1a377c3b324564f24a9021da84d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/__init__.py @@ -0,0 +1,66 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING + +from .utils.import_utils import _LazyModule + +if TYPE_CHECKING: + from .agent_template import BaseAgentTemplate, agent_template_map + from .arguments import (AppArguments, BaseArguments, DeployArguments, EvalArguments, ExportArguments, + InferArguments, PretrainArguments, RLHFArguments, RolloutArguments, SamplingArguments, + SftArguments) + from .callbacks import TrainerCallback, callbacks_map + from .dataset import EncodePreprocessor, load_dataset + from .infer_engine import (AdapterRequest, GRPOVllmEngine, InferClient, InferEngine, InferRequest, LmdeployEngine, + RequestConfig, SglangEngine, TransformersEngine, VllmEngine) + from .loss import BaseLoss, loss_map + from .loss_scale import ALL_BASE_STRATEGY, ConfigLossScale, LossScale, get_loss_scale, loss_scale_map + from .metrics import InferStats, MeanMetric, eval_metrics_map + from .model import get_model_processor, get_processor + from .optimizers import OptimizerCallback, optimizers_map + from .pipelines import (app_main, deploy_main, eval_main, export_main, infer_main, merge_lora, pretrain_main, + rlhf_main, rollout_main, run_deploy, sampling_main, sft_main) + from .template import get_template + from .trainers import Seq2SeqTrainer, Seq2SeqTrainingArguments, Trainer, TrainingArguments + from .tuner_plugin import PeftTuner, Tuner, tuners_map + from .tuners import Swift + from .utils import get_logger, safe_snapshot_download + from .version import __release_datetime__, __version__ +else: + _import_structure = { + 'version': ['__release_datetime__', '__version__'], + 'tuners': ['Swift'], + 'tuner_plugin': ['Tuner', 'PeftTuner', 'tuners_map'], + 'infer_engine': [ + 'TransformersEngine', 'VllmEngine', 'SglangEngine', 'LmdeployEngine', 'InferRequest', 'RequestConfig', + 'AdapterRequest', 'InferEngine', 'InferClient', 'GRPOVllmEngine' + ], + 'trainers': ['TrainingArguments', 'Seq2SeqTrainingArguments', 'Trainer', 'Seq2SeqTrainer'], + 'arguments': [ + 'PretrainArguments', 'SftArguments', 'RLHFArguments', 'ExportArguments', 'InferArguments', 'AppArguments', + 'EvalArguments', 'SamplingArguments', 'RolloutArguments', 'DeployArguments', 'BaseArguments' + ], + 'pipelines': [ + 'sft_main', 'pretrain_main', 'infer_main', 'rlhf_main', 'export_main', 'app_main', 'eval_main', + 'sampling_main', 'rollout_main', 'deploy_main', 'merge_lora', 'run_deploy' + ], + 'model': ['get_model_processor', 'get_processor'], + 'template': ['get_template'], + 'dataset': ['load_dataset', 'EncodePreprocessor'], + 'utils': ['get_logger', 'safe_snapshot_download'], + 'agent_template': ['agent_template_map', 'BaseAgentTemplate'], + 'loss': ['loss_map', 'BaseLoss'], + 'metrics': ['eval_metrics_map', 'InferStats', 'MeanMetric'], + 'optimizers': ['optimizers_map', 'OptimizerCallback'], + 'callbacks': ['callbacks_map', 'TrainerCallback'], + 'loss_scale': ['loss_scale_map', 'LossScale', 'get_loss_scale', 'ALL_BASE_STRATEGY', 'ConfigLossScale'], + } + + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..90a3edbe2736d51f295361dc095d7fc3d0a9ebfd --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import BaseAgentTemplate +from .mapping import agent_template_map diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/base.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/base.py new file mode 100644 index 0000000000000000000000000000000000000000..57fcc302e52d0b3b42721e383ff43f155ffb96fa --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/base.py @@ -0,0 +1,261 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +""" +Agent template module for handling tool calling and function execution. + +This module provides base classes and utilities for creating agent templates +that support tool calling in conversational AI systems. It includes support +for various agent formats like ReAct, function calling, and parallel execution. +""" +import ast +import json +from abc import ABC, abstractmethod +from dataclasses import asdict, dataclass +from typing import Any, Dict, List, Literal, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt, split_str_parts_by + + +@dataclass +class AgentKeyword: + action: str = 'Action:' + action_input: str = 'Action Input:' + observation: str = 'Observation:' + + +@dataclass +class ToolDesc: + name_for_model: str + name_for_human: str + description_for_model: str + parameters: str + args_format: str + + +class ReactCompatMixin: + """ + Mixin class providing ReAct-style agent compatibility. + + This mixin handles parsing and formatting of tool calls in the ReAct format, + where actions and inputs are marked with specific keywords in the text. + """ + keyword = AgentKeyword() + + @staticmethod + def _split_action_action_input(response: str, keyword: AgentKeyword) -> List[Function]: + agent_parts = split_str_parts_by(response, list(asdict(keyword).values())) + functions = [] + action_content = None + + for part in agent_parts: + key, content = part['key'].lower(), part['content'] + if action_content is None and key == keyword.action.lower(): + action_content = content + elif action_content is not None and key == keyword.action_input.lower(): + functions.append(Function(name=action_content, arguments=content)) + action_content = None + + return functions + + def get_toolcall(self, response: str) -> List[Function]: + """ + Extract tool calls from an agent response. + + Args: + response: The agent's response text. + + Returns: + List of Function objects representing tool calls. + """ + functions = self._split_action_action_input(response, self.keyword) + if len(functions) == 0 and self.keyword != ReactCompatMixin.keyword: + # compat react + functions = self._split_action_action_input(response, ReactCompatMixin.keyword) + return functions + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + """ + Format tool execution results into the conversation. + + Args: + assistant_content: The assistant's message containing tool calls. + tool_messages: List of tool execution result messages. + + Returns: + Tuple of (formatted assistant content, formatted tool responses). + """ + assert len(tool_messages) > 0 + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + if not assistant_content.endswith(self.keyword.observation): + if not assistant_content.endswith('\n'): + assistant_content += '\n' + assistant_content += self.keyword.observation + res = [] + for i, tool_message in enumerate(tool_messages): + if i > 0: + res.append(self.keyword.observation) + tool_content = tool_message['content'] + res.append(tool_content) + if not tool_content.endswith('\n'): + res.append('\n') + else: + res = [] + for tool_message in tool_messages: + res.append(tool_message['content']) + return assistant_content, res + + def _format_react_standalone_tool_responses(self, tool_messages) -> 'Prompt': + assert len(tool_messages) > 0 + res = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res.extend(['\n', self.keyword.observation, '\n', tool_content]) + return res + + def _format_standalone_tool_responses(self, tool_messages) -> 'Prompt': + """Format observations that do not follow an assistant tool call.""" + raise NotImplementedError( + f'Agent template {self.__class__.__name__} does not support standalone tool responses.') + + @staticmethod + def _parse_tool_call(content) -> Dict[str, Any]: + obj = BaseAgentTemplate._parse_json(content) + name = obj['name'] + arguments = obj.get('arguments') + if arguments is None: + arguments = obj.get('parameters') + arguments = BaseAgentTemplate._parse_json(arguments) + assert arguments is not None, f'content: {content}' + return {'name': name, 'arguments': arguments} + + def _format_tool_calls(self, tool_call_messages) -> str: + """ + Format tool call messages into ReAct format. + + Args: + tool_call_messages: List of messages containing tool call information. + + Returns: + Formatted string with Action, Action Input, and Observation markers. + """ + # -> assistant_content + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + tool_calls.append(f'{self.keyword.action} {tool_call["name"]}\n' + f'{self.keyword.action_input} {tool_call["arguments"]}\n') + tool_calls.append(self.keyword.observation) + return ''.join(tool_calls) + + +class BaseAgentTemplate(ReactCompatMixin, ABC): + """ + Abstract base class for agent templates. + + This class provides common functionality for parsing and formatting tools, + as well as handling tool calls in different formats. Subclasses must + implement the following methods to define their specific behavior: + - `_format_tools`: Format tool definitions for the prompt + - `_format_tool_calls`: Format tool call messages + - `_format_tool_responses`: Format tool execution results + - `get_toolcall`: Extract tool calls from agent responses + """ + + def _add_tool_call_prefix(self, tool_content: str, pre_message=None) -> str: + """Hook to prepend a separator before tool_call content based on the preceding message. + + Subclasses can override this to match their jinja template's separator logic + (e.g., Qwen3.5/3.6 inserts '\n\n' when assistant has effective content before tool_calls). + + Args: + tool_content: The formatted tool_call string from _format_tool_calls. + pre_message: The message immediately before the tool_call block, or None. + + Returns: + tool_content with any necessary prefix prepended. + """ + return tool_content + + @staticmethod + def _get_tool_name(tool): + return tool.get('name_for_model') or tool.get('name') + + @staticmethod + def unwrap_tool(tool): + assert isinstance(tool, dict), f'tool: {tool}' + if 'type' in tool and 'function' in tool: + tool = tool['function'] + return tool + + @staticmethod + def wrap_tool(tool): + assert isinstance(tool, dict), f'tool: {tool}' + if 'type' not in tool and 'function' not in tool: + tool = {'type': 'function', 'function': tool} + return tool + + @staticmethod + def _parse_tool(tool, lang: Literal['zh', 'en']) -> ToolDesc: + tool = BaseAgentTemplate.unwrap_tool(tool) + name_for_model = BaseAgentTemplate._get_tool_name(tool) + name_for_human = tool.get('name_for_human') or name_for_model + + description = tool.get('description') + if description is None: + description = tool.get('description_for_model') + parameters = tool.get('parameters') or {} + parameters = parameters if isinstance(parameters, str) else json.dumps(parameters, ensure_ascii=False) + args_format = '此工具的输入应为JSON对象。' if lang == 'zh' else 'Format the arguments as a JSON object.' + tool_desc = ToolDesc( + name_for_model=name_for_model, + name_for_human=name_for_human, + description_for_model=description, + parameters=parameters, + args_format=args_format) + assert name_for_model is not None and description is not None, f'tool_desc: {tool_desc}' + return tool_desc + + @staticmethod + def _parse_json(json_str: str) -> Optional[Any]: + """ + Parse a JSON string with fallback to ast.literal_eval. + + Args: + json_str: String to parse, or already parsed object. + + Returns: + Parsed object, or None if parsing fails. + """ + if not isinstance(json_str, str): + return json_str + try: + res = json.loads(json_str) + except json.JSONDecodeError: + try: + res = ast.literal_eval(json_str) + except Exception: + return + return res + + @abstractmethod + def _format_tools(self, + tools: List[Union[str, dict]], + system: Optional[str] = None, + user_message: Optional[dict] = None) -> str: + """ + Format tools for inclusion in the agent prompt. + + Args: + tools: List of tool definitions (strings or dictionaries). + system: System prompt text. + user_message: Optional user message to incorporate. + + Returns: + Formatted string to include in the prompt. + """ + pass diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/deepseek_v3_1.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/deepseek_v3_1.py new file mode 100644 index 0000000000000000000000000000000000000000..f468d33c7d053181b91fd717387188e33773bf8b --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/deepseek_v3_1.py @@ -0,0 +1,87 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class DeepSeekV31AgentTemplate(BaseAgentTemplate): + + def get_toolcall(self, response: str) -> List[Function]: + # Parse tool calls using the DSV3.1 format: + # <｜tool▁calls▁begin｜><｜tool▁call▁begin｜>name<｜tool▁sep｜>args<｜tool▁call▁end｜> + pattern = r'<｜tool▁call▁begin｜>(.*?)<｜tool▁sep｜>(.*?)<｜tool▁call▁end｜>' + res_list = re.findall(pattern, response, re.DOTALL) + functions = [] + for name, arguments in res_list: + name = name.strip() + arguments = self._parse_json(arguments.strip()) + if arguments is not None: + functions.append(Function(name=name, arguments=arguments)) + + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _get_tool_responses(self, tool_messages): + return ''.join(f'<｜tool▁output▁begin｜>{tool_message["content"]}<｜tool▁output▁end｜>' + for tool_message in tool_messages) + + def _get_tool_calls(self, tool_calls: List[str]): + return f'<｜tool▁calls▁begin｜>{"".join(tool_calls)}<｜tool▁calls▁end｜>' + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + res = ['<｜end▁of▁sentence｜>', self._get_tool_responses(tool_messages)] + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [] + system = system or '' + for tool in tools: + tool = self.unwrap_tool(tool) + tool_name = self._get_tool_name(tool) + description = tool.get('description', '') + parameters = tool.get('parameters', {}) + + tool_desc = f"""### {tool_name} +Description: {description} + +Parameters: {json.dumps(parameters, ensure_ascii=False)}""" + tool_descs.append(tool_desc) + + tools_section = '\n\n'.join(tool_descs) + + return f"""{system} + +## Tools +You have access to the following tools: + +{tools_section} + +IMPORTANT: ALWAYS adhere to this exact format for tool use: +<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>tool_call_name<｜tool▁sep｜>tool_call_arguments<｜tool▁call▁end｜>{{additional_tool_calls}}<｜tool▁calls▁end｜> + +Where: +- `tool_call_name` must be an exact match to one of the available tools +- `tool_call_arguments` must be valid JSON that strictly follows the tool's Parameters Schema +- For multiple tool calls, chain them directly without separators or spaces""" + + def _format_tool_calls(self, tool_call_messages): + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + name = tool_call['name'] + arguments = json.dumps(tool_call['arguments'], ensure_ascii=False) + tool_calls.append(f'<｜tool▁call▁begin｜>{name}<｜tool▁sep｜>{arguments}<｜tool▁call▁end｜>') + return self._get_tool_calls(tool_calls) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/deepseek_v4.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/deepseek_v4.py new file mode 100644 index 0000000000000000000000000000000000000000..6f787641d7be24e4659cf8dff05b38e15e375a98 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/deepseek_v4.py @@ -0,0 +1,145 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import Any, Dict, List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + +DSML_TOKEN = '｜DSML｜' + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. \ +You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. \ +For all other types (numbers, booleans, arrays, objects), \ +pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by ), \ +you MUST output your complete reasoning inside ... BEFORE any tool calls or final response. + +Otherwise, output directly after with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + + +def _to_json(value: Any) -> str: + try: + return json.dumps(value, ensure_ascii=False) + except Exception: + return json.dumps(value, ensure_ascii=True) + + +def _encode_arguments_to_dsml(arguments: Dict[str, Any]) -> str: + """Encode tool call arguments dict into DSML parameter lines.""" + lines = [] + for k, v in arguments.items(): + is_str = 'true' if isinstance(v, str) else 'false' + val = v if isinstance(v, str) else _to_json(v) + lines.append(f'<{DSML_TOKEN}parameter name="{k}" string="{is_str}">{val}') + return '\n'.join(lines) + + +class DeepSeekV4AgentTemplate(BaseAgentTemplate): + + def get_toolcall(self, response: str) -> List[Function]: + # Parse DSML tool calls from model output + # Pattern: <｜DSML｜invoke name="tool_name">...params... + invoke_pattern = re.compile( + rf'<{re.escape(DSML_TOKEN)}invoke\s+name="([^"]+)">\s*(.*?)\s*', re.DOTALL) + param_pattern = re.compile( + rf'<{re.escape(DSML_TOKEN)}parameter\s+name="([^"]+)"\s+string="(true|false)">' + rf'(.*?)', re.DOTALL) + + functions = [] + for match in invoke_pattern.finditer(response): + tool_name = match.group(1) + params_block = match.group(2) + arguments = {} + for pm in param_pattern.finditer(params_block): + param_name = pm.group(1) + is_string = pm.group(2) + param_value = pm.group(3) + if is_string == 'false': + try: + param_value = json.loads(param_value) + except json.JSONDecodeError: + pass + arguments[param_name] = param_value + functions.append(Function(name=tool_name, arguments=json.dumps(arguments, ensure_ascii=False))) + + if len(functions) == 0: + # Fallback to ReAct format + return super().get_toolcall(response) + return functions + + def _get_tool_responses(self, tool_messages): + # The official encoding merges tool results into one user turn, joining the + # `` blocks with a blank line. + return '\n\n'.join(f'{tool_message["content"]}' for tool_message in tool_messages) + + def _add_tool_call_prefix(self, tool_content: str, pre_message=None) -> str: + # The official encoding always renders the tool_calls block as `\n\n` + block, + # right after the assistant's (possibly empty) textual content. + return '\n\n' + tool_content + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + res = [ + '<｜end▁of▁sentence｜><｜User｜>', + self._get_tool_responses(tool_messages), + '<｜Assistant｜>', + ] + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_schemas = [] + for tool in tools: + tool = self.unwrap_tool(tool) + tool_schemas.append(_to_json(tool)) + + tools_section = TOOLS_TEMPLATE.format( + tool_schemas='\n'.join(tool_schemas), + dsml_token=DSML_TOKEN, + ) + + system = system or '' + return f'{system}\n\n{tools_section}' if system else tools_section + + def _format_tool_calls(self, tool_call_messages) -> str: + invocations = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + name = tool_call['name'] + arguments = tool_call['arguments'] + if isinstance(arguments, str): + arguments = json.loads(arguments) + dsml_args = _encode_arguments_to_dsml(arguments) + invocations.append(f'<{DSML_TOKEN}invoke name="{name}">\n{dsml_args}\n') + + tool_calls_str = '\n'.join(invocations) + return f'<{DSML_TOKEN}tool_calls>\n{tool_calls_str}\n' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/extra.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/extra.py new file mode 100644 index 0000000000000000000000000000000000000000..f325f647660ede84e463bd3d71fa71d01b88da3d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/extra.py @@ -0,0 +1,36 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import List, Optional, Union + +from .base import BaseAgentTemplate + + +class ReactGRPOAgentTemplate(BaseAgentTemplate): + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_names = [] + tool_descs = [] + for tool in tools: + tool_desc = self._parse_tool(tool, 'en') + tool_names.append(tool_desc.name_for_model) + tool_descs.append( + f'{tool_desc.name_for_model}: Call this tool to interact with the {tool_desc.name_for_human} API. ' + f'What is the {tool_desc.name_for_human} API useful for? {tool_desc.description_for_model} ' + f'Parameters: {tool_desc.parameters} {tool_desc.args_format}') + + return """A conversation for tool calling between User and Assistant. The user asks a question which may be solved by calling tools, and the Assistant solves it. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process should be enclosed within tags and answer should follow the ReACT format(Action:xxx\nAction Input:xxx), i.e., reasoning process here Action: action here\nAction Input: parameters here + +Answer the following questions as best as you can. You have access to the following tools: + +""" + '\n\n'.join(tool_descs) + f""" + +Use the following format: + +you should always think about what to do +Action: the action to take, should be one of [{','.join(tool_names)}] +Action Input: the input to the action +Observation: the result of the action, given by the actual calling +... (this Thought/Action/Action Input/Observation can be repeated zero or more times) +Final Answer: the final answer to the original input question + +Begin! +""" # noqa diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/gemma4.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/gemma4.py new file mode 100644 index 0000000000000000000000000000000000000000..3ffcd76efa72057a345a3537dbbf3e55c507b8ca --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/gemma4.py @@ -0,0 +1,210 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import Any, Dict, List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + +QUOTE = '<|"|>' +_STANDARD_KEYS = {'description', 'type', 'properties', 'required', 'nullable'} + + +class Gemma4AgentTemplate(BaseAgentTemplate): + """Agent template for Google Gemma-4 models. + + Reference: chat_template.jinja shipped with google/gemma-4-12B-it. + Tool definitions are wrapped in `<|tool>...` and rendered with the + custom DSL described by the official chat template. + Tool calls follow `<|tool_call>call:NAME{key:value,...}` and + tool responses follow `<|tool_response>response:NAME{...}`. + """ + + @classmethod + def _format_argument(cls, value: Any, escape_keys: bool = True) -> str: + if isinstance(value, bool): + return 'true' if value else 'false' + if isinstance(value, str): + return f'{QUOTE}{value}{QUOTE}' + if value is None: + return 'null' + if isinstance(value, dict): + items = [] + for k in sorted(value.keys()): + v = value[k] + key_str = f'{QUOTE}{k}{QUOTE}' if escape_keys else str(k) + items.append(f'{key_str}:{cls._format_argument(v, escape_keys=escape_keys)}') + return '{' + ','.join(items) + '}' + if isinstance(value, (list, tuple)): + return '[' + ','.join(cls._format_argument(item, escape_keys=escape_keys) for item in value) + ']' + return str(value) + + @classmethod + def _format_parameters(cls, + properties: Dict[str, Any], + required: Optional[List[str]] = None, + filter_keys: bool = False) -> str: + parts = [] + for key in sorted(properties.keys()): + value = properties[key] + if filter_keys and key in _STANDARD_KEYS: + continue + if not isinstance(value, dict): + continue + inner: List[str] = [] + type_upper = (value.get('type') or '').upper() if isinstance(value.get('type'), str) else '' + if value.get('description'): + inner.append(f'description:{QUOTE}{value["description"]}{QUOTE}') + if type_upper == 'STRING': + if value.get('enum'): + inner.append(f'enum:{cls._format_argument(value["enum"])}') + elif type_upper == 'ARRAY': + items_value = value.get('items') + if isinstance(items_value, dict) and items_value: + items_inner: List[str] = [] + items_required = items_value.get('required', []) + for item_key in sorted(items_value.keys()): + item_value = items_value[item_key] + if item_value is None: + continue + if item_key == 'properties' and isinstance(item_value, dict): + items_inner.append(f'properties:{{{cls._format_parameters(item_value, items_required)}}}') + elif item_key == 'required': + req_str = ','.join(f'{QUOTE}{r}{QUOTE}' for r in item_value) + items_inner.append(f'required:[{req_str}]') + elif item_key == 'type': + if isinstance(item_value, str): + items_inner.append(f'type:{cls._format_argument(item_value.upper())}') + else: + items_inner.append(f'type:{cls._format_argument([str(t).upper() for t in item_value])}') + else: + items_inner.append(f'{item_key}:{cls._format_argument(item_value)}') + inner.append('items:{' + ','.join(items_inner) + '}') + if value.get('nullable'): + inner.append('nullable:true') + if type_upper == 'OBJECT': + inner_required = value.get('required', []) + if isinstance(value.get('properties'), dict): + inner.append(f'properties:{{{cls._format_parameters(value["properties"], inner_required)}}}') + else: + inner.append(f'properties:{{{cls._format_parameters(value, inner_required, filter_keys=True)}}}') + if value.get('required'): + req_str = ','.join(f'{QUOTE}{r}{QUOTE}' for r in value['required']) + inner.append(f'required:[{req_str}]') + inner.append(f'type:{QUOTE}{type_upper}{QUOTE}') + parts.append(f'{key}:{{{",".join(inner)}}}') + return ','.join(parts) + + @classmethod + def _format_function_declaration(cls, tool: Dict[str, Any]) -> str: + function = tool['function'] + name = function.get('name', '') + description = function.get('description', '') or '' + result = f'declaration:{name}{{description:{QUOTE}{description}{QUOTE}' + params = function.get('parameters') + if params: + param_parts: List[str] = [] + properties = params.get('properties') + if properties: + param_parts.append(f'properties:{{{cls._format_parameters(properties, params.get("required", []))}}}') + if params.get('required'): + req_str = ','.join(f'{QUOTE}{r}{QUOTE}' for r in params['required']) + param_parts.append(f'required:[{req_str}]') + ptype = params.get('type') + if isinstance(ptype, str) and ptype: + param_parts.append(f'type:{QUOTE}{ptype.upper()}{QUOTE}') + if param_parts: + result += ',parameters:{' + ','.join(param_parts) + '}' + result += '}' + return result + + def _format_tools(self, + tools: List[Union[str, dict]], + system: Optional[str] = None, + user_message: Optional[dict] = None) -> str: + tool_blocks: List[str] = [] + for tool in tools: + tool = self.wrap_tool(tool) + tool_blocks.append(f'<|tool>{self._format_function_declaration(tool)}') + system_text = (system or '').strip() + return system_text + ''.join(tool_blocks) + + def _format_tool_calls(self, tool_call_messages) -> str: + invocations: List[str] = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + name = tool_call['name'] + arguments = tool_call['arguments'] + if isinstance(arguments, str): + arguments = self._parse_json(arguments) or {} + if isinstance(arguments, dict): + args_str = ','.join(f'{k}:{self._format_argument(arguments[k], escape_keys=False)}' + for k in sorted(arguments.keys())) + else: + args_str = '' + invocations.append(f'<|tool_call>call:{name}{{{args_str}}}') + return ''.join(invocations) + + def _get_tool_responses(self, tool_messages) -> str: + parts: List[str] = [] + for tool_message in tool_messages: + tool_name = tool_message.get('name') or 'unknown' + tool_content = tool_message.get('content') + if isinstance(tool_content, dict): + inner = ','.join(f'{k}:{self._format_argument(tool_content[k], escape_keys=False)}' + for k in sorted(tool_content.keys())) + parts.append(f'<|tool_response>response:{tool_name}{{{inner}}}') + else: + # Match jinja: treat string/other content as a single `value:` field. + value = '' if tool_content is None else tool_content + parts.append(f'<|tool_response>response:{tool_name}' + f'{{value:{self._format_argument(value, escape_keys=False)}}}') + return ''.join(parts) + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + # If the model hallucinated a trailing `<|tool_response>` opener (e.g. when stop + # tokens were not configured), strip it so the rendered turn does not contain + # `<|tool_response><|tool_response>response:...`. + if assistant_content.endswith('<|tool_response>'): + assistant_content = assistant_content[:-len('<|tool_response>')] + # In gemma4, tool_call/tool_response/follow-up assistant text all live in the + # same `<|turn>model ... ` block, so we do not open a new model turn here. + res: 'Prompt' = [self._get_tool_responses(tool_messages)] + return assistant_content, res + + @classmethod + def _gemma_to_json(cls, s: str) -> str: + # `<|"|>` -> `"`; bare keys preceded by `{` or `,` get JSON-quoted. + s = s.replace(QUOTE, '"') + s = re.sub(r'(?<=[\{,])([A-Za-z_][\w\-]*)(?=:)', r'"\1"', s) + return s + + @classmethod + def _parse_arguments(cls, args_body: str) -> Dict[str, Any]: + json_str = cls._gemma_to_json('{' + args_body + '}') + try: + parsed = json.loads(json_str) + if isinstance(parsed, dict): + return parsed + except json.JSONDecodeError: + pass + return {} + + def get_toolcall(self, response: str) -> List[Function]: + pattern = re.compile(r'<\|tool_call>call:([^\{]+)\{(.*?)\}', re.DOTALL) + functions: List[Function] = [] + for match in pattern.finditer(response): + name = match.group(1).strip() + arguments = self._parse_arguments(match.group(2)) + functions.append(Function(name=name, arguments=json.dumps(arguments, ensure_ascii=False))) + if not functions: + return super().get_toolcall(response) + return functions diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/glm4.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/glm4.py new file mode 100644 index 0000000000000000000000000000000000000000..7b3844947e9294a2a2109e78255a22ed66de8a3c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/glm4.py @@ -0,0 +1,194 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class ChatGLM4AgentTemplate(BaseAgentTemplate): + is_glm4_0414 = False + + @staticmethod + def _find_function_call(single_content: str) -> Optional[Function]: + single_content = single_content.replace('<|observation|>', '') + pattern = re.compile(r'([^\n`]*?)\n({.*?})(?=\w*\n|$)', re.DOTALL) + matches = pattern.findall(single_content) + if not matches: + return + name, arguments = matches[0] + return Function(name=name, arguments=arguments) + + def get_toolcall(self, response: str) -> List[Function]: + toolcall_list = response.split('<|assistant|>') + functions = [] + for toolcall in toolcall_list: + function = self._find_function_call(toolcall) + if function: + functions.append(function) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [] + for tool in tools: + tool = self.unwrap_tool(tool) + name = self._get_tool_name(tool) + tool_descs.append(f'## {name}\n\n{json.dumps(tool, ensure_ascii=False, indent=4)}\n' + '在调用上述函数时，请使用 Json 格式表示调用的参数。') + glm4_system = '你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。\n\n' # noqa + return ('' if self.is_glm4_0414 else glm4_system) + """# 可用工具 + +""" + '\n'.join(tool_descs) + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + res = ['\n'] + for i, tool_message in enumerate(tool_messages): + tool_content = tool_message['content'] + if i > 0: + res.append('<|observation|>\n') + res.append(tool_content) + res.append('<|assistant|>\n') + return assistant_content, res + + def _format_standalone_tool_responses(self, tool_messages) -> 'Prompt': + res = [] + for tool_message in tool_messages: + res.extend(['<|observation|>\n', tool_message['content']]) + return res + + def _format_tool_calls(self, tool_call_messages) -> str: + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + tool_calls.append(f'{tool_call["name"]}\n{tool_call["arguments"]}') + return '<|assistant|>'.join(tool_calls) + '<|observation|>' + + +class GLM4AgentTemplate(ChatGLM4AgentTemplate): + is_glm4_0414 = True + + +class GLM4_5AgentTemplate(BaseAgentTemplate): + model_type = 'glm4_5' + + @staticmethod + def _find_function_call(single_content: str) -> Optional[Function]: + single_content = single_content.strip() + func_name_match = re.match(r'^([^\n<]+)', single_content) + if not func_name_match: + return None + func_name = func_name_match.group(1).strip() + keys = re.findall(r'(.*?)', single_content, re.DOTALL) + values = re.findall(r'(.*?)', single_content, re.DOTALL) + if len(keys) != len(values): + return None + args = {k.strip(): v.strip() for k, v in zip(keys, values)} + return Function(name=func_name, arguments=json.dumps(args, ensure_ascii=False)) + + def get_toolcall(self, response: str) -> List[Function]: + toolcall_list = re.findall(r'(.*?)', response, re.DOTALL) + functions = [] + for toolcall in toolcall_list: + function = self._find_function_call(toolcall) + if function: + functions.append(function) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [ + '# Tools\n\nYou may call one or more functions to assist with the user query.\n\n' + 'You are provided with function signatures within XML tags:\n' + ] + for tool in tools: + if self.model_type == 'glm5_1': + tool = self.unwrap_tool(tool) + tool_descs.append(f'{json.dumps(tool, ensure_ascii=False)}') + if self.model_type == 'glm4_5': + tool_desc = ('\n\nFor each function call, output the function name and arguments within ' + 'the following XML format:\n{function-name}\n{arg-key-1}\n' + '{arg-value-1}\n{arg-key-2}\n' + '{arg-value-2}\n...\n') + elif self.model_type in {'glm4_7', 'glm5_1'}: + tool_desc = ('\n\nFor each function call, output the function name and arguments within ' + 'the following XML format:\n{function-name}{arg-key-1}' + '{arg-value-1}{arg-key-2}' + '{arg-value-2}...') + else: + raise ValueError("model_type must be one of 'glm4_5', 'glm4_7', or 'glm5_1'.") + tool_descs.append(tool_desc) + tool_descs = '\n'.join(tool_descs) + if system is not None and system.strip(): + tool_descs += '<|system|>\n' + system.strip() + elif self.model_type in {'glm4_7', 'glm5_1'} and not tool_descs.startswith('\n'): + tool_descs = '\n' + tool_descs + return tool_descs + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + if self.model_type == 'glm4_5': + res = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res.append(f'\n\n{tool_content}\n') + res.append('<|assistant|>\n') + elif self.model_type in {'glm4_7', 'glm5_1'}: + res = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res.append(f'{tool_content}') + res.append('<|assistant|>') + return assistant_content, res + + def _format_standalone_tool_responses(self, tool_messages) -> 'Prompt': + res = ['<|observation|>'] + for tool_message in tool_messages: + tool_content = tool_message['content'] + if self.model_type == 'glm4_5': + res.append(f'\n\n{tool_content}\n') + elif self.model_type in {'glm4_7', 'glm5_1'}: + res.append(f'{tool_content}') + return res + + def _format_tool_calls(self, tool_call_messages) -> str: + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + tool_calls.append(f"{tool_call['name']}") + for arg_key, arg_value in tool_call['arguments'].items(): + tool_calls.append(f'{arg_key}') + tool_calls.append(f'{arg_value}') + tool_calls.append('') + if self.model_type == 'glm4_5': + sep = '\n' + elif self.model_type in {'glm4_7', 'glm5_1'}: + sep = '' + return sep.join(tool_calls) + '<|observation|>' + + +class GLM4_7AgentTemplate(GLM4_5AgentTemplate): + model_type = 'glm4_7' + + +class GLM5_1AgentTemplate(GLM4_5AgentTemplate): + model_type = 'glm5_1' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/hermes.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/hermes.py new file mode 100644 index 0000000000000000000000000000000000000000..79071fb691e1ca5a7550aa9d780c43d104a3e2e0 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/hermes.py @@ -0,0 +1,144 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class HermesAgentTemplate(BaseAgentTemplate): + + def get_toolcall(self, response: str) -> List[Function]: + res_list = re.findall(r'(.+?)', response, re.DOTALL) + functions = [] + for res in res_list: + res = self._parse_json(res) + if isinstance(res, dict) and 'name' in res and 'arguments' in res: + functions.append(Function(name=res['name'], arguments=res['arguments'])) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _get_tool_responses(self, tool_messages): + res_tool = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res_tool.append(f'\n{tool_content}\n') + return '\n'.join(res_tool) + + def _get_tool_calls(self, tool_calls: List[str]): + return '\n'.join(tool_calls) + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + if hasattr(self, 'template_meta'): + prompt = self.template_meta.prompt + chat_sep = self.template_meta.chat_sep + else: + prompt = ['<|im_start|>user\n{{QUERY}}<|im_end|>\n<|im_start|>assistant\n'] + chat_sep = ['<|im_end|>\n'] + res = chat_sep.copy() + total_tool = self._get_tool_responses(tool_messages) + for context in prompt: + if isinstance(context, str): + context = context.replace('{{QUERY}}', total_tool) + res.append(context) + return assistant_content, res + + def _format_standalone_tool_responses(self, tool_messages) -> 'Prompt': + """Render standalone results as a native tool-response user turn.""" + if not hasattr(self, 'template_meta'): + return super()._format_standalone_tool_responses(tool_messages) + + res = (self.template_meta.chat_sep or []).copy() + total_tool = self._get_tool_responses(tool_messages) + for context in self.template_meta.prompt: + if isinstance(context, str) and '{{QUERY}}' in context: + query_prefix = context.split('{{QUERY}}', maxsplit=1)[0] + if query_prefix: + res.append(query_prefix) + res.append(total_tool) + return res + res.append(context) + raise ValueError(f'Template prompt does not contain {{{{QUERY}}}}: {self.template_meta.prompt}') + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools] + system = system or '' + return f"""{system} + +# Tools + +You may call one or more functions to assist with the user query. + +You are provided with function signatures within XML tags: + +""" + '\n'.join(tool_descs) + """ + + +For each function call, return a json object with function name and arguments within XML tags: + +{"name": , "arguments": } +""" + + def _format_tool_calls(self, tool_call_messages): + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + tool_calls.append(f'\n{json.dumps(tool_call, ensure_ascii=False)}\n') + return self._get_tool_calls(tool_calls) + + +class HunyuanHermesAgentTemplate(HermesAgentTemplate): + + def get_toolcall(self, response: str) -> List[Function]: + res_list = re.findall(r'(.+?)\n```json(.+?)```', response, re.DOTALL) + functions = [] + for name, arguments in res_list: + arguments = self._parse_json(arguments) + functions.append(Function(name=name, arguments=arguments)) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _get_tool_responses(self, tool_messages): + res_tool = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res_tool.append(f'{tool_content}') + tool_responses = '\n'.join(res_tool) + return f'{tool_responses}' + + def _get_tool_calls(self, tool_calls: List[str]): + tool_calls = '\n'.join(tool_calls) + return f'\n{tool_calls}\n' + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools] + system = system or '' + if system: + system = f'{system}\n\n' + return f"""{system}# Tools + +You may call one or more functions to assist with the user query. + +You are provided with function signatures within XML tags: + +""" + '\n'.join(tool_descs) + """ + + +For function call returns, you should first print For each function call, you should return object like: +function_name +```json +function_arguments_in_json_format +```At the end of function call returns, you should print """ diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/hy_v3.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/hy_v3.py new file mode 100644 index 0000000000000000000000000000000000000000..c9d52d0c29f8ad65f6a06fedccea6695ba5ec641 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/hy_v3.py @@ -0,0 +1,107 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class HyV3PreviewAgentTemplate(BaseAgentTemplate): + HYTK = '' + + def get_toolcall(self, response: str) -> List[Function]: + # Parse tool calls from ...name............ + tool_call_blocks = re.findall(rf'(.*?)', response, re.DOTALL) + functions = [] + for block in tool_call_blocks: + # Extract function name: text before + name_match = re.match(rf'(.*?)', block, re.DOTALL) + if not name_match: + continue + name = name_match.group(1).strip() + # Extract arg_key/arg_value pairs together to avoid misalignment + pairs = re.findall( + rf'(.*?)\s*(.*?)', + block, re.DOTALL) + arguments = {} + for k, v in pairs: + k = k.strip() + v = v.strip() + parsed = self._parse_json(v) + arguments[k] = parsed if parsed is not None else v + functions.append(Function(name=name, arguments=arguments)) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _get_tool_responses(self, tool_messages): + res_tool = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res_tool.append(f'\n{tool_content}\n') + tool_responses = '\n'.join(res_tool) + return f'\n{tool_responses}\n' + + def _get_tool_calls(self, tool_calls: List[str]): + tool_calls_str = '\n'.join(tool_calls) + return f'\n{tool_calls_str}\n' + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + res = [f'<｜hy_eos{self.HYTK}｜>', self._get_tool_responses(tool_messages), f'<｜hy_Assistant{self.HYTK}｜>'] + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools] + system = system or '' + if system: + system = f'{system}\n\n' + return f"""{system}# Tools + +You may call one or more functions to assist with the user query. + +You are provided with function signatures within XML tags: + +""" + '\n'.join(tool_descs) + f""" + + +For function call returns, you should first print +For each function call, you should return object like: +{{function-name}} +{{arg-key-1}} +{{arg-value-1}} +{{arg-key-2}} +{{arg-value-2}} +... + +At the end of function call returns, you should print """ + + def _format_tool_calls(self, tool_call_messages): + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + name = tool_call['name'] + arguments = tool_call['arguments'] + arg_lines = [] + if isinstance(arguments, dict): + for k, v in arguments.items(): + if not isinstance(v, str): + v = json.dumps(v, ensure_ascii=False) + arg_lines.append(f'{k}\n' + f'{v}') + arg_str = '\n'.join(arg_lines) + tool_calls.append(f'{name}\n{arg_str}\n') + return self._get_tool_calls(tool_calls) + + +class HyV3AgentTemplate(HyV3PreviewAgentTemplate): + HYTK = ':opensource' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/kimi_k25.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/kimi_k25.py new file mode 100644 index 0000000000000000000000000000000000000000..17c1e7fad251894f939e6e8a9a9dcba3c3a9ed47 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/kimi_k25.py @@ -0,0 +1,119 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class KimiK25AgentTemplate(BaseAgentTemplate): + """Agent template for Kimi K2.5/K2.6 models. + + Tool calling format: + - Tools are declared in a separate system message with role 'tool_declare' + using TypeScript namespace format. + - Tool calls: + <|tool_calls_section_begin|> + <|tool_call_begin|>{function_name}<|tool_call_argument_begin|>{args_json}<|tool_call_end|> + <|tool_calls_section_end|> + - Tool response: + <|im_system|>tool<|im_middle|>## Return of {tool_call_id} + {content}<|im_end|> + """ + + @staticmethod + def _json_type_to_ts(json_type): + type_map = { + 'string': 'string', + 'number': 'number', + 'integer': 'number', + 'boolean': 'boolean', + 'array': 'any[]', + 'object': 'object', + 'null': 'null', + } + return type_map.get(json_type, 'any') + + def _tools_to_typescript(self, tools): + parts = [] + for tool in tools: + tool = self.unwrap_tool(tool) + name = self._get_tool_name(tool) + description = tool.get('description', '') + parameters = tool.get('parameters', {}) + properties = parameters.get('properties', {}) + + lines = [] + if description: + lines.append(f'// {description}') + + if not properties: + lines.append(f'type {name} = (_: {{}}) => any;') + else: + lines.append(f'type {name} = (_: {{') + props = list(properties.items()) + for i, (pname, pschema) in enumerate(props): + pdesc = pschema.get('description', '') + ptype = self._json_type_to_ts(pschema.get('type', '')) + if pdesc: + lines.append(f' // {pdesc}') + if i < len(props) - 1: + lines.append(f' {pname}: {ptype},') + else: + lines.append(f' {pname}: {ptype}') + lines.append('}) => any;') + parts.append('\n'.join(lines)) + return '\n'.join(parts) + + def get_toolcall(self, response: str) -> List[Function]: + pattern = r'<\|tool_call_begin\|>(.*?)<\|tool_call_argument_begin\|>(.*?)<\|tool_call_end\|>' + res_list = re.findall(pattern, response, re.DOTALL) + functions = [] + for name, arguments in res_list: + name = name.strip() + arguments = arguments.strip() + parsed_args = self._parse_json(arguments) + if parsed_args is not None: + functions.append(Function(name=name, arguments=parsed_args)) + else: + functions.append(Function(name=name, arguments=arguments)) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + ts_tools = self._tools_to_typescript(tools) + tool_content = f'# Tools\n\n## functions\nnamespace functions {{\n{ts_tools}\n}}\n' + system = system or '' + res = f'tool_declare<|im_middle|>{tool_content}' + if system: + res += f'<|im_end|><|im_system|>system<|im_middle|>{system}' + return res + + def _format_tool_calls(self, tool_call_messages) -> str: + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + name = tool_call['name'] + arguments = json.dumps(tool_call['arguments'], ensure_ascii=False) + tool_calls.append(f'<|tool_call_begin|>{name}<|tool_call_argument_begin|>{arguments}<|tool_call_end|>') + return f'<|tool_calls_section_begin|>{"".join(tool_calls)}<|tool_calls_section_end|>' + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + res = ['<|im_end|>'] + for tool_message in tool_messages: + tool_call_id = tool_message.get('tool_call_id', '') + tool_content = tool_message['content'] + res.append(f'<|im_system|>tool<|im_middle|>## Return of {tool_call_id}\n{tool_content}<|im_end|>') + res.append('<|im_assistant|>assistant<|im_middle|>') + return assistant_content, res diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/kimi_k3.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/kimi_k3.py new file mode 100644 index 0000000000000000000000000000000000000000..e59866b44567c2309b12f466c193fedac177c964 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/kimi_k3.py @@ -0,0 +1,151 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class KimiK3AgentTemplate(BaseAgentTemplate): + """Agent template for Kimi K3 models (XTML format). + + Aligned with `encoding_k3.py` in the model repo: + - Tools are declared in a separate system message with `type="tool-declare"`, + carrying the deep-sorted JSONSchema in a compact ```json block. + - Tool calls follow the closed response channel: + <|close|>response<|sep|><|open|>tools<|sep|> + <|open|>call tool="{name}" index="{i}"<|sep|> + <|open|>argument key="{key}" type="{type}"<|sep|>{value}<|close|>argument<|sep|> + <|close|>call<|sep|><|close|>tools<|sep|> + - Tool response: + <|open|>message role="tool" tool="{name}" index="{i}"<|sep|>{content}<|close|>message<|sep|><|end_of_msg|> + """ + + @staticmethod + def _deep_sort(obj): + if isinstance(obj, dict): + return {k: KimiK3AgentTemplate._deep_sort(v) for k, v in sorted(obj.items())} + if isinstance(obj, list): + return [KimiK3AgentTemplate._deep_sort(item) for item in obj] + return obj + + @staticmethod + def _escape_attr(value) -> str: + return str(value).replace('&', '&').replace('"', '"') + + @staticmethod + def _unescape_attr(value: str) -> str: + return value.replace('"', '"').replace('&', '&') + + @staticmethod + def _xtml_type(value) -> str: + if isinstance(value, bool): + return 'boolean' + if value is None: + return 'null' + if isinstance(value, (int, float)): + return 'number' + if isinstance(value, str): + return 'string' + if isinstance(value, dict): + return 'object' + return 'array' + + @staticmethod + def _xtml_value(value) -> str: + if isinstance(value, str): + return value + return json.dumps(value, ensure_ascii=False) + + def get_toolcall(self, response: str) -> List[Function]: + call_pattern = r'<\|open\|>call tool="(.*?)"(?: index="\d+")?<\|sep\|>(.*?)<\|close\|>call<\|sep\|>' + arg_pattern = r'<\|open\|>argument key="(.*?)" type="(.*?)"<\|sep\|>(.*?)<\|close\|>argument<\|sep\|>' + json_pattern = r'<\|open\|>json type="object"<\|sep\|>(.*?)<\|close\|>json<\|sep\|>' + functions = [] + for name, body in re.findall(call_pattern, response, re.DOTALL): + name = self._unescape_attr(name) + json_match = re.search(json_pattern, body, re.DOTALL) + if json_match: + arguments = self._parse_json(json_match.group(1)) + if arguments is None: + arguments = json_match.group(1) + else: + arguments = {} + for key, arg_type, value in re.findall(arg_pattern, body, re.DOTALL): + key = self._unescape_attr(key) + if arg_type == 'string': + arguments[key] = value + elif arg_type == 'null': + arguments[key] = None + else: + parsed = self._parse_json(value) + arguments[key] = value if parsed is None else parsed + functions.append(Function(name=name, arguments=arguments)) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tools = self._deep_sort([self.wrap_tool(tool) for tool in tools]) + tools_json = json.dumps(tools, ensure_ascii=False, separators=(',', ':')) + tool_content = ('# Tools\nHere are the available tools, described in JSONSchema.\n\n' + f'```json\n{tools_json}\n```') + # Hijack the `<|open|>message {{SYSTEM}}<|close|>message<|sep|><|end_of_msg|>` system + # prefix so the tool-declare message is rendered before the plain system message. + res = f'role="system" type="tool-declare"<|sep|>{tool_content}' + if system: + res += f'<|close|>message<|sep|><|end_of_msg|><|open|>message role="system"<|sep|>{system}' + return res + + def _format_tool_calls(self, tool_call_messages) -> str: + calls = [] + for index, message in enumerate(tool_call_messages, start=1): + tool_call = self._parse_tool_call(message['content']) + name, arguments = tool_call['name'], tool_call['arguments'] + call = [f'<|open|>call tool="{self._escape_attr(name)}" index="{index}"<|sep|>'] + if isinstance(arguments, dict): + for key, value in arguments.items(): + call.append(f'<|open|>argument key="{self._escape_attr(key)}" ' + f'type="{self._xtml_type(value)}"<|sep|>' + f'{self._xtml_value(value)}<|close|>argument<|sep|>') + else: + # non-object arguments fall back to a raw json block + json_block = arguments if isinstance(arguments, str) else json.dumps(arguments, ensure_ascii=False) + call.append(f'<|open|>json type="object"<|sep|>{json_block}<|close|>json<|sep|>') + call.append('<|close|>call<|sep|>') + calls.append(''.join(call)) + return f'<|open|>tools<|sep|>{"".join(calls)}<|close|>tools<|sep|>' + + def _add_tool_call_prefix(self, tool_content: str, pre_message=None) -> str: + # Without a preceding assistant text message (whose converted content already ends + # with `<|close|>response<|sep|>`), the empty response channel must be closed here. + if pre_message is None or pre_message['role'] != 'assistant': + tool_content = '<|close|>response<|sep|>' + tool_content + return tool_content + + def _format_tool_responses( + self, + assistant_content, + tool_messages, + ) -> Tuple[str, 'Prompt']: + contents = assistant_content if isinstance(assistant_content, list) else [assistant_content] + joined = ''.join(c for c in contents if isinstance(c, str)) + with_action = self.keyword.action in joined and self.keyword.action_input in joined + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + # `tool`/`index` attributes are resolved from the preceding assistant tool calls, + # in call order (aligned with the `tool_index` enumeration in encoding_k3.py). + names = re.findall(r'<\|open\|>call tool="(.*?)" index="\d+"<\|sep\|>', joined) + res = ['<|close|>message<|sep|><|end_of_msg|>'] + for index, tool_message in enumerate(tool_messages, start=1): + if index <= len(names): + name = self._unescape_attr(names[index - 1]) + else: + name = tool_message.get('tool') or tool_message.get('name') or '' + res.append(f'<|open|>message role="tool" tool="{self._escape_attr(name)}" index="{index}"<|sep|>' + f'{tool_message["content"]}<|close|>message<|sep|><|end_of_msg|>') + res.append('<|open|>message role="assistant"<|sep|>') + return assistant_content, res diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/llama.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/llama.py new file mode 100644 index 0000000000000000000000000000000000000000..e12cbd18e7078f266cf7c7fb3b101bbec39cae50 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/llama.py @@ -0,0 +1,74 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class Llama3AgentTemplate(BaseAgentTemplate): + eom_token = '<|eom_id|>' + start_token = '<|start_header_id|>' + end_token = '<|end_header_id|>' + eot_token = '<|eot_id|>' + + def get_toolcall(self, response: str) -> List[Function]: + if response.endswith(self.eom_token): + response = response[:-len(self.eom_token)] + functions = [] + res_list = re.findall(r'{[^{]*?"name":.*?"parameters":\s*?{.*?}\s*?}', response, re.DOTALL) + for res in res_list: + res = self._parse_json(res) + if isinstance(res, dict) and 'name' in res and 'parameters' in res: + functions.append(Function(name=res['name'], arguments=res['parameters'])) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + res = [self.eot_token] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res.append(f'{self.start_token}tool{self.end_token}\n\n{tool_content}{self.eot_token}') + res.append(f'{self.start_token}assistant{self.end_token}\n\n') + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + assert user_message is not None + user_content = user_message['content'] + tool_descs = [json.dumps(tool, ensure_ascii=False, indent=4) for tool in tools] + new_user_content = """Given the following functions, please respond with a JSON for a function call with its proper arguments that best answers the given prompt. + +Respond in the format {"name": function name, "parameters": dictionary of argument name and its value}. Do not use variables. + +""" + '\n\n'.join(tool_descs) + f""" + +{user_content}""" # noqa + user_message['content'] = new_user_content + return system or '' + + def _format_tool_calls(self, tool_call_messages) -> str: + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + tool_call['parameters'] = tool_call.pop('arguments') + tool_calls.append(json.dumps(tool_call, ensure_ascii=False)) + return '\n'.join(tool_calls) + + +class Llama4AgentTemplate(Llama3AgentTemplate): + eom_token = '<|eom|>' + start_token = '<|header_start|>' + end_token = '<|header_end|>' + eot_token = '<|eot|>' + toolcall_pattern = r'(.+?)<\|eom\|>' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/mapping.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/mapping.py new file mode 100644 index 0000000000000000000000000000000000000000..9609b2a35a545bc80d00773e12a580c760e5c03d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/mapping.py @@ -0,0 +1,63 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .deepseek_v3_1 import DeepSeekV31AgentTemplate +from .deepseek_v4 import DeepSeekV4AgentTemplate +from .extra import ReactGRPOAgentTemplate +from .gemma4 import Gemma4AgentTemplate +from .glm4 import (ChatGLM4AgentTemplate, GLM4_5AgentTemplate, GLM4_7AgentTemplate, GLM4AgentTemplate, + GLM5_1AgentTemplate) +from .hermes import HermesAgentTemplate, HunyuanHermesAgentTemplate +from .hy_v3 import HyV3AgentTemplate, HyV3PreviewAgentTemplate +from .kimi_k3 import KimiK3AgentTemplate +from .kimi_k25 import KimiK25AgentTemplate +from .llama import Llama3AgentTemplate, Llama4AgentTemplate +from .minicpm5 import MiniCPM5AgentTemplate +from .minimax_m2 import MinimaxM2AgentTemplate +from .minimax_m3 import MinimaxM3AgentTemplate +from .mistral import MistralAgentTemplate +from .qwen import QwenEnAgentTemplate, QwenEnParallelAgentTemplate, QwenZhAgentTemplate, QwenZhParallelAgentTemplate +from .qwen3_coder import Qwen3_5AgentTemplate, Qwen3CoderAgentTemplate +from .react import ReactEnAgentTemplate, ReactZnAgentTemplate +from .seed_oss import SeedAgentTemplate +from .toolbench import ToolBenchAgentTemplate +from .youtu import YoutuAgentTemplate + +agent_template_map = { + # ref: https://qwen.readthedocs.io/zh-cn/latest/framework/function_call.html#function-calling-templates + 'react_en': ReactEnAgentTemplate, + 'react_zh': ReactZnAgentTemplate, + # ref: https://github.com/QwenLM/Qwen-Agent/blob/main/qwen_agent/llm/fncall_prompts/qwen_fncall_prompt.py + 'qwen_en': QwenEnAgentTemplate, + 'qwen_zh': QwenZhAgentTemplate, + 'qwen_en_parallel': QwenEnParallelAgentTemplate, + 'qwen_zh_parallel': QwenZhParallelAgentTemplate, + 'qwen3_coder': Qwen3CoderAgentTemplate, + 'qwen3_5': Qwen3_5AgentTemplate, + 'hermes': HermesAgentTemplate, + 'hunyuan_hermes': HunyuanHermesAgentTemplate, + 'hy_v3_preview': HyV3PreviewAgentTemplate, + 'hy_v3': HyV3AgentTemplate, + 'toolbench': ToolBenchAgentTemplate, # ref: https://modelscope.cn/datasets/swift/ToolBench + 'chatglm4': ChatGLM4AgentTemplate, + 'glm4': GLM4AgentTemplate, # ref: https://modelscope.cn/models/ZhipuAI/GLM-4-9B-0414 + 'glm4_5': GLM4_5AgentTemplate, + 'glm4_7': GLM4_7AgentTemplate, + 'glm5_1': GLM5_1AgentTemplate, + 'llama3': Llama3AgentTemplate, + 'llama4': Llama4AgentTemplate, + # ref: https://huggingface.co/deepseek-ai/DeepSeek-V3.1 + 'deepseek_v3_1': DeepSeekV31AgentTemplate, + # ref: https://modelscope.cn/models/deepseek-ai/DeepSeek-V4-Flash + 'deepseek_v4': DeepSeekV4AgentTemplate, + 'minimax_m2': MinimaxM2AgentTemplate, + 'minimax_m3': MinimaxM3AgentTemplate, + 'seed_oss': SeedAgentTemplate, + # ref: https://modelscope.cn/models/google/gemma-4-12B-it + 'gemma4': Gemma4AgentTemplate, + # extra + 'react_grpo': ReactGRPOAgentTemplate, + 'mistral': MistralAgentTemplate, + 'youtu': YoutuAgentTemplate, + 'kimi_k25': KimiK25AgentTemplate, + 'kimi_k3': KimiK3AgentTemplate, + 'minicpm5': MiniCPM5AgentTemplate, +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minicpm5.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minicpm5.py new file mode 100644 index 0000000000000000000000000000000000000000..98f490d59ee915d63c07d4b017821f9664ad2004 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minicpm5.py @@ -0,0 +1,114 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class MiniCPM5AgentTemplate(BaseAgentTemplate): + """Agent template for MiniCPM5 models using XML-based function calling format. + + Tool call format: + param-value + + Tool response format: + + response_content + + """ + + def get_toolcall(self, response: str) -> List[Function]: + # Match ... blocks + func_pattern = re.compile(r'(.*?)', re.DOTALL) + param_pattern = re.compile(r'' + r'(?:|([^<]*))' + r'', re.DOTALL) + + functions = [] + for func_match in func_pattern.finditer(response): + func_name = func_match.group(1) + func_body = func_match.group(2) + arguments = {} + for param_match in param_pattern.finditer(func_body): + param_name = param_match.group(1) + # CDATA value or plain value + param_value = param_match.group(2) if param_match.group(2) is not None else param_match.group(3) + # Try to parse as JSON value (number, bool, etc.) + try: + param_value = json.loads(param_value) + except (json.JSONDecodeError, ValueError): + pass + arguments[param_name] = param_value + functions.append(Function(name=func_name, arguments=arguments)) + + if len(functions) == 0: + # Fallback to ReAct-style parsing + return super().get_toolcall(response) + return functions + + def _get_tool_responses(self, tool_messages): + res_tool = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res_tool.append(f'\n{tool_content}\n') + return '\n'.join(res_tool) + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + if hasattr(self, 'template_meta'): + prompt = self.template_meta.prompt + chat_sep = self.template_meta.chat_sep + else: + prompt = ['<|im_start|>user\n{{QUERY}}<|im_end|>\n<|im_start|>assistant\n'] + chat_sep = ['<|im_end|>\n'] + res = chat_sep.copy() + total_tool = self._get_tool_responses(tool_messages) + for context in prompt: + if isinstance(context, str): + context = context.replace('{{QUERY}}', total_tool) + res.append(context) + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools] + system = system or '' + if system: + system = f'{system}\n\n' + return (f'{system}# Tools\n\n' + 'You are provided with function signatures within XML tags:\n' + '\n' + '\n'.join(tool_descs) + '\n\n\n' + 'Tool usage guidelines:\n' + '- You may call zero or more functions. If no function calls are needed, ' + 'just answer normally and do not include any .\n' + '- When calling a function, return an XML object within using:\n' + 'param-value\n' + '- param-value may be multi-line. If it contains <, & or newline characters, ' + 'wrap it in a CDATA block: ') + + def _format_tool_calls(self, tool_call_messages) -> str: + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + name = tool_call['name'] + arguments = tool_call['arguments'] + params_xml = '' + if isinstance(arguments, dict): + for param_name, param_value in arguments.items(): + value_str = param_value if isinstance(param_value, str) else json.dumps( + param_value, ensure_ascii=False) + if isinstance(param_value, str) and ('<' in param_value or '&' in param_value + or '\n' in param_value): + params_xml += f'' + else: + params_xml += f'{value_str}' + tool_calls.append(f'{params_xml}') + return '\n'.join(tool_calls) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minimax_m2.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minimax_m2.py new file mode 100644 index 0000000000000000000000000000000000000000..628ef7b69aa44a75762146aadff01492ee3d98fe --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minimax_m2.py @@ -0,0 +1,171 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class MinimaxM2AgentTemplate(BaseAgentTemplate): + """ + Agent template for MiniMax-M2 series models. + + This template handles tool calling in MiniMax's XML-based format: + + + param-value + + + """ + + def get_toolcall(self, response: str) -> List[Function]: + """ + Extract tool calls from MiniMax response format. + + Format: + + + param-value + + + """ + functions = [] + + # Find all tool_call blocks + tool_call_blocks = re.findall(r'(.*?)', response, re.DOTALL) + + for block in tool_call_blocks: + # Find all invoke blocks within the tool_call + invoke_blocks = re.findall(r'(.*?)', block, re.DOTALL) + + for tool_name, params_block in invoke_blocks: + # Extract parameters + params = {} + param_matches = re.findall(r'(.*?)', params_block, re.DOTALL) + + for param_name, param_value in param_matches: + param_value = param_value.strip() + # Try to parse as JSON if it looks like a JSON structure + parsed_value = self._parse_json(param_value) + params[param_name] = parsed_value if parsed_value is not None else param_value + + functions.append(Function(name=tool_name, arguments=params)) + + # Fallback to react format if no functions found + if len(functions) == 0: + return super().get_toolcall(response) + + return functions + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + """ + Format tool execution results in MiniMax format. + + Tool responses are wrapped in tags. + """ + # Check if using react format + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + + # Use template meta if available + if hasattr(self, 'template_meta'): + prompt = self.template_meta.prompt.copy() + chat_sep = self.template_meta.chat_sep + for i in range(len(prompt)): + if isinstance(prompt[i], str): + prompt[i] = prompt[i].replace('user', 'tool') + else: + # Default format based on the Jinja2 template + prompt = [']~b]tool\n{{QUERY}}[e~[\n'] + chat_sep = ['[e~[\n'] + + res = chat_sep.copy() + + # Format tool responses + tool_responses = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + tool_responses.append(f'{tool_content}') + + total_tool = '\n'.join(tool_responses) + + for context in prompt: + if isinstance(context, str): + context = context.replace('{{QUERY}}', total_tool) + res.append(context) + + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + """ + Format tools in MiniMax format with JSONSchema and XML invocation examples. + """ + # Parse tools to JSONSchema format + tool_schemas = [] + for tool in tools: + tool = self.unwrap_tool(tool) + tool_schemas.append(json.dumps(tool, ensure_ascii=False)) + + system = system or '' + + return f"""{system} + +# Tools +You may call one or more tools to assist with the user query. +Here are the tools available in JSONSchema format: + + +""" + '\n'.join(f'{schema}' for schema in tool_schemas) + """ + + +When making tool calls, use XML format to invoke tools and pass parameters: + + + +param-value-1 +param-value-2 +... + +""" + + def _format_tool_calls(self, tool_call_messages): + """ + Format tool call messages into MiniMax XML format. + + Args: + tool_call_messages: List of messages containing tool call information. + + Returns: + Formatted string with tool calls in MiniMax XML format. + """ + tool_calls = [] + + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + name = tool_call['name'] + arguments = tool_call['arguments'] + + # Build parameter list + params = [] + for key, value in arguments.items(): + # Convert value to JSON string if it's not a string + if not isinstance(value, str): + value = json.dumps(value, ensure_ascii=False) + params.append(f'{value}') + + # Build invoke block + invoke_block = f'\n' + '\n'.join(params) + '\n' + tool_calls.append(invoke_block) + + # Wrap all invocations in tool_call tags + if tool_calls: + return '\n' + '\n'.join(tool_calls) + '\n' + + return '' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minimax_m3.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minimax_m3.py new file mode 100644 index 0000000000000000000000000000000000000000..4188bec8011a8486c42c74c8b6d2d2bc3ad5e4dc --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/minimax_m3.py @@ -0,0 +1,235 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import Any, List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + +# Special token used as a namespace prefix for every XML tag in MiniMax-M3 +# tool_call payloads. +NS_TOKEN = ']<]minimax[>[' +TOOLCALL_BEGIN_TOKEN = NS_TOKEN + '' +TOOLCALL_END_TOKEN = NS_TOKEN + '' + + +def _to_xml(val: Any, ns: str = NS_TOKEN) -> str: + """Recursive XML renderer mirroring the ``to_xml`` macro in MiniMax-M3's + ``chat_template.jinja``. + + ``None`` values are intentionally omitted (consistent with the upstream + convention that drops ``None`` parameters rather than emitting a literal + ``null`` string). + """ + if val is None: + return '' + if isinstance(val, dict): + parts = [] + for k, v in val.items(): + if v is None: + continue + parts.append(f'{ns}<{k}>{_to_xml(v, ns)}{ns}') + return ''.join(parts) + if isinstance(val, (list, tuple)): + parts = [] + for item in val: + parts.append(f'{ns}{_to_xml(item, ns)}{ns}') + return ''.join(parts) + if isinstance(val, bool): + return json.dumps(val) + return str(val) + + +_NS = re.escape(NS_TOKEN) +_TC_BEGIN = re.escape(TOOLCALL_BEGIN_TOKEN) +_TC_END = re.escape(TOOLCALL_END_TOKEN) +# Match any opening tag like ]<]minimax[>[ (excluding closing/invoke/tool_call) +_INVOKE_RE = re.compile(rf'{_NS}(.*?){_NS}', re.DOTALL) +_TOOLCALL_RE = re.compile(rf'{_TC_BEGIN}(.*?){_TC_END}', re.DOTALL) + + +def _parse_xml_value(content: str) -> Any: + """Parse an XML fragment produced by ``to_xml`` back into a Python value. + + The expected fragments use ``NS_TOKEN`` as a tag prefix. The function + handles nested ```` lists, dict-like ``...`` structures + and falls back to a stripped string for primitive payloads. + """ + content = content.strip() + if not content: + return '' + + # Try list of items first (heuristic: starts with ``). + if content.startswith(f'{NS_TOKEN}'): + items = [] + for inner in _iter_tagged(content, 'item'): + items.append(_parse_xml_value(inner)) + return items + + # Try mapping (heuristic: starts with a NS_TOKEN). + if content.startswith(NS_TOKEN + '<'): + result: dict = {} + for key, inner in _iter_keyed(content): + result[key] = _parse_xml_value(inner) + if result: + return result + + # Primitive fallback. Try JSON (booleans / numbers) before raw text. + try: + return json.loads(content) + except Exception: + return content + + +def _iter_tagged(content: str, tag: str): + pattern = re.compile(rf'{_NS}<{re.escape(tag)}>(.*?){_NS}', re.DOTALL) + for m in pattern.finditer(content): + yield m.group(1) + + +def _iter_keyed(content: str): + """Iterate ``(tag_name, inner_content)`` for top-level NS-prefixed tags.""" + cursor = 0 + n = len(content) + open_pat = re.compile(rf'{_NS}<([^/!?\s>]+)>') + while cursor < n: + m = open_pat.search(content, cursor) + if not m: + return + name = m.group(1) + end_marker = f'{NS_TOKEN}' + # Match nested same-name tags by counting depth. + depth = 1 + scan = m.end() + open_marker = f'{NS_TOKEN}<{name}>' + while depth > 0 and scan < n: + next_open = content.find(open_marker, scan) + next_close = content.find(end_marker, scan) + if next_close == -1: + return + if next_open != -1 and next_open < next_close: + depth += 1 + scan = next_open + len(open_marker) + else: + depth -= 1 + scan = next_close + len(end_marker) + inner = content[m.end():scan - len(end_marker)] + yield name, inner + cursor = scan + + +class MinimaxM3AgentTemplate(BaseAgentTemplate): + """Agent template for MiniMax-M3 series multimodal models. + + Tool calls follow this XML-with-namespace format: + + ]<]minimax[>[ + ]<]minimax[>[ + ]<]minimax[>[value-1]<]minimax[>[ + ]<]minimax[>[]<]minimax[>[...]<]minimax[>[]<]minimax[>[ + ]<]minimax[>[ + ]<]minimax[>[ + + Tool responses are wrapped in ``...`` inside a + ``]~b]tool`` slot. + """ + + def get_toolcall(self, response: str) -> List[Function]: + functions: List[Function] = [] + for tc_block in _TOOLCALL_RE.findall(response): + for tool_name, params_block in _INVOKE_RE.findall(tc_block): + arguments = {} + for key, inner in _iter_keyed(params_block): + arguments[key] = _parse_xml_value(inner) + functions.append(Function(name=tool_name, arguments=arguments)) + + if not functions: + return super().get_toolcall(response) + return functions + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + + if hasattr(self, 'template_meta'): + prompt = self.template_meta.prompt.copy() + chat_sep = self.template_meta.chat_sep + for i in range(len(prompt)): + if isinstance(prompt[i], str): + prompt[i] = prompt[i].replace('user', 'tool') + else: + prompt = [']~b]tool\n{{QUERY}}[e~[\n]~b]ai\n'] + chat_sep = ['[e~[\n'] + + res = chat_sep.copy() if chat_sep else [] + tool_responses = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + tool_responses.append(f'{tool_content}') + total_tool = '\n'.join(tool_responses) + + for context in prompt: + if isinstance(context, str): + context = context.replace('{{QUERY}}', total_tool) + res.append(context) + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_schemas = [] + for tool in tools: + tool = self.unwrap_tool(tool) + tool_schemas.append(json.dumps(tool, ensure_ascii=False)) + + system = system or '' + tools_xml = '\n'.join(f'{schema}' for schema in tool_schemas) + # Mirror the example block produced by chat_template.jinja so the + # in-context format hint matches inference time exactly. + # Note: jinja emits 'Example:\n' then '\n' before the tool_call_begin + # token, which renders as two consecutive newlines. + example = (f'\n\n{TOOLCALL_BEGIN_TOKEN}\n' + f'{NS_TOKEN}' + f'{NS_TOKEN}value-1{NS_TOKEN}' + f'{NS_TOKEN}' + f'{NS_TOKEN}' + f'{NS_TOKEN}val-a{NS_TOKEN}' + f'{NS_TOKEN}val-b{NS_TOKEN}' + f'{NS_TOKEN}' + f'{NS_TOKEN}' + f'{NS_TOKEN}\n' + f'{NS_TOKEN}' + f'{NS_TOKEN}value-1{NS_TOKEN}' + f'{NS_TOKEN}\n' + f'{TOOLCALL_END_TOKEN}') + + return (f'{system}\n\n# Tools\n' + 'You may call one or more tools to assist with the user query.\n' + 'Here are the tools available in JSONSchema format:\n' + f'\n\n{tools_xml}\n\n\n' + f'To call tools, wrap all invocations in a single {TOOLCALL_BEGIN_TOKEN}{TOOLCALL_END_TOKEN} ' + 'block. Parameter values containing nested objects or arrays are recursively expanded into ' + f'XML elements. Example:{example}') + + def _format_tool_calls(self, tool_call_messages) -> str: + invocations = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + name = tool_call['name'] + arguments = tool_call['arguments'] or {} + + param_parts = [f'{NS_TOKEN}'] + for k, v in arguments.items(): + if v is None: + continue + param_parts.append(f'{NS_TOKEN}<{k}>{_to_xml(v, NS_TOKEN)}{NS_TOKEN}') + param_parts.append(f'{NS_TOKEN}') + invocations.append(''.join(param_parts)) + + if not invocations: + return '' + return f'{TOOLCALL_BEGIN_TOKEN}\n' + '\n'.join(invocations) + f'\n{TOOLCALL_END_TOKEN}' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/mistral.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/mistral.py new file mode 100644 index 0000000000000000000000000000000000000000..adf649e6896b0e162d5111ccb5771be0be55dcb4 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/mistral.py @@ -0,0 +1,67 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class MistralAgentTemplate(BaseAgentTemplate): + + def get_toolcall(self, response: str) -> List[Function]: + res_list = re.findall(r'\[TOOL_CALLS\]\[(.*?)\]', response, re.DOTALL) + if not res_list: + return [] + res_list = res_list[0].strip().split('\n') + functions = [] + for res_str in res_list: + parsed_res = self._parse_json(res_str) + if isinstance(parsed_res, dict): + parsed_res = [parsed_res] # Handle single tool call + if isinstance(parsed_res, list): + for tool_call in parsed_res: + if isinstance(tool_call, dict) and 'name' in tool_call and 'arguments' in tool_call: + functions.append(Function(name=tool_call['name'], arguments=tool_call['arguments'])) + if len(functions) == 0: + # compat react_en + return super().get_toolcall(response) + return functions + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + if not hasattr(self, 'template_meta'): + raise ValueError('MistralAgentTemplate requires template_meta to be registered') + prompt = self.template_meta.prompt + chat_sep = self.template_meta.chat_sep + + res = chat_sep.copy() + res_tool = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + # append `[TOOL_RESULTS]{"content": {{ .Content }}}[/TOOL_RESULTS]` to res_tool + res_tool.append(f'[TOOL_RESULTS]{json.dumps({"content": tool_content}, ensure_ascii=False)}[/TOOL_RESULTS]') + total_tool = '\n'.join(res_tool) + for context in prompt: + if isinstance(context, str): + context = context.replace('{{QUERY}}', total_tool) + res.append(context) + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools] + system = system or '' + return f"""{system}[AVAILABLE_TOOLS]{' '.join(tool_descs)}[/AVAILABLE_TOOLS]""" + + def _format_tool_calls(self, tool_call_messages): + tool_calls = [] + for message in tool_call_messages: + # needs `{'name': name, 'arguments': arguments}`, which self._parse_tool_call + # satisfies + tool_call = self._parse_tool_call(message['content']) + tool_calls.append(json.dumps(tool_call, ensure_ascii=False)) + return f'[TOOL_CALLS][\n{chr(10).join(tool_calls)}\n]' # check if need `` at end diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/qwen.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/qwen.py new file mode 100644 index 0000000000000000000000000000000000000000..dddd0a6ec6444cadbbffbebe755a3a712b5436bc --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/qwen.py @@ -0,0 +1,134 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import List, Optional, Union + +from .base import AgentKeyword, BaseAgentTemplate + +keyword = AgentKeyword( + action='✿FUNCTION✿:', + action_input='✿ARGS✿:', + observation='✿RESULT✿:', +) + + +class QwenEnAgentTemplate(BaseAgentTemplate): + keyword = keyword + + def _get_tool_names_descs(self, tools): + tool_names = [] + tool_descs = [] + for tool in tools: + tool_desc = self._parse_tool(tool, 'en') + tool_names.append(tool_desc.name_for_model) + tool_descs.append(f'### {tool_desc.name_for_human}\n\n' + f'{tool_desc.name_for_model}: {tool_desc.description_for_model} ' + f'Parameters: {tool_desc.parameters} {tool_desc.args_format}') + return tool_names, tool_descs + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_names, tool_descs = self._get_tool_names_descs(tools) + system = system or '' + return f"""{system} + +# Tools + +## You have access to the following tools: + +""" + '\n\n'.join(tool_descs) + f""" + +## When you need to call a tool, please insert the following command in your reply, which can be called zero or multiple times according to your needs: + +✿FUNCTION✿: The tool to use, should be one of [{','.join(tool_names)}] +✿ARGS✿: The input of the tool +✿RESULT✿: Tool results +✿RETURN✿: Reply based on tool results. Images need to be rendered as ![](url)""" # noqa + + +class QwenZhAgentTemplate(BaseAgentTemplate): + keyword = keyword + + def _get_tool_names_descs(self, tools): + tool_names = [] + tool_descs = [] + for tool in tools: + tool_desc = self._parse_tool(tool, 'zh') + tool_names.append(tool_desc.name_for_model) + tool_descs.append(f'### {tool_desc.name_for_human}\n\n' + f'{tool_desc.name_for_model}: {tool_desc.description_for_model} ' + f'输入参数：{tool_desc.parameters} {tool_desc.args_format}') + return tool_names, tool_descs + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_names, tool_descs = self._get_tool_names_descs(tools) + system = system or '' + return f"""{system} + +# 工具 + +## 你拥有如下工具： + +""" + '\n\n'.join(tool_descs) + f""" + +## 你可以在回复中插入零次、一次或多次以下命令以调用工具： + +✿FUNCTION✿: 工具名称，必须是[{','.join(tool_names)}]之一。 +✿ARGS✿: 工具输入 +✿RESULT✿: 工具结果 +✿RETURN✿: 根据工具结果进行回复，需将图片用![](url)渲染出来""" # noqa + + +class QwenEnParallelAgentTemplate(QwenEnAgentTemplate): + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_names, tool_descs = self._get_tool_names_descs(tools) + system = system or '' + return f"""{system} + +# Tools + +## You have access to the following tools: + +""" + '\n\n'.join(tool_descs) + f""" + +## Insert the following command in your reply when you need to call N tools in parallel: + +✿FUNCTION✿: The name of tool 1, should be one of [{','.join(tool_names)}] +✿ARGS✿: The input of tool 1 +✿FUNCTION✿: The name of tool 2 +✿ARGS✿: The input of tool 2 +... +✿FUNCTION✿: The name of tool N +✿ARGS✿: The input of tool N +✿RESULT✿: The result of tool 1 +✿RESULT✿: The result of tool 2 +... +✿RESULT✿: he result of tool N +✿RETURN✿: Reply based on tool results. Images need to be rendered as ![](url)""" # noqa + + +class QwenZhParallelAgentTemplate(QwenZhAgentTemplate): + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_names, tool_descs = self._get_tool_names_descs(tools) + system = system or '' + return f"""{system} + +# 工具 + +## 你拥有如下工具： + +""" + '\n\n'.join(tool_descs) + f""" + +## 你可以在回复中插入以下命令以并行调用N个工具： + +✿FUNCTION✿: 工具1的名称，必须是[{','.join(tool_names)}]之一 +✿ARGS✿: 工具1的输入 +✿FUNCTION✿: 工具2的名称 +✿ARGS✿: 工具2的输入 +... +✿FUNCTION✿: 工具N的名称 +✿ARGS✿: 工具N的输入 +✿RESULT✿: 工具1的结果 +✿RESULT✿: 工具2的结果 +... +✿RESULT✿: 工具N的结果 +✿RETURN✿: 根据工具结果进行回复，需将图片用![](url)渲染出来""" # noqa diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/qwen3_coder.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/qwen3_coder.py new file mode 100644 index 0000000000000000000000000000000000000000..5c260d7b10d84467fb6ed812926487afef8cfeb5 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/qwen3_coder.py @@ -0,0 +1,200 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Union + +from swift.infer_engine import Function +from .hermes import HermesAgentTemplate + + +def render_extra_keys(obj, handled_keys): + """Helper function to render extra keys not explicitly handled""" + result = '' + if isinstance(obj, dict): + for key, value in obj.items(): + if key not in handled_keys: + result += f'\n<{key}>{json.dumps(value, ensure_ascii=False)}' + return result + + +TOOL_DESC_SUFFIX = ( + '\n\nIf you choose to call a function ONLY reply in the following format with ' + 'NO suffix:\n\n\n\n\n' + 'value_1\n\n\nThis is the value for the second parameter\n' + 'that can span\nmultiple lines\n\n\n\n\n\n' + 'Reminder:\n- Function calls MUST follow the specified format: an inner ' + 'block must be nested within XML tags\n- Required parameters MUST be specified\n' + '- You may provide optional reasoning for your function call in natural language BEFORE the function call, ' + 'but NOT after\n- If there is no function call available, ' + 'answer the question like normal with your current ' + 'knowledge and do not tell the user about function calls\n') + + +class Qwen3CoderAgentTemplate(HermesAgentTemplate): + + @staticmethod + def _find_function_call(single_content: str) -> Optional[Function]: + single_content = single_content.strip() + # Check whether the complete function tag is included + if not single_content.startswith(''): + return None + + # Extract function name + func_name_match = re.search(r']+)>', single_content) + if not func_name_match: + return None + + func_name = func_name_match.group(1).strip() + parameters = {} + + # Use regular expressions to match parameters + # Match any content of content + param_pattern = r']+)>\s*(.*?)\s*' + param_matches = re.findall(param_pattern, single_content, re.DOTALL) + + for param_name, param_value in param_matches: + # Clear the parameter values and remove any possible additional whitespace + clean_value = param_value.strip() + parameters[param_name.strip()] = clean_value + + return Function(name=func_name, arguments=json.dumps(parameters, ensure_ascii=False)) + + def get_toolcall(self, response: str) -> List[Function]: + # Extract the tool call parameters from the model's response + toolcall_list = re.findall(r'(.*?)', response, re.DOTALL) + functions = [] + for toolcall in toolcall_list: + function = self._find_function_call(toolcall) + if function: + functions.append(function) + if len(functions) == 0: + # Compat react_en + return super(HermesAgentTemplate, self).get_toolcall(response) + return functions + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + if system is None: + system = 'You are Qwen, a helpful AI assistant that can interact with a computer to solve tasks.' + tool_descs = [f'{system}\n\n# Tools\n\nYou have access to the following functions:\n\n'] + for tool in tools: + tool_desc = '' + + # Check function key + if isinstance(tool, dict) and 'function' in tool: + tool = tool['function'] + + # Add function name + tool_desc += f"\n{tool['name']}" + + # Add description if available + if 'description' in tool: + tool_desc += f"\n{tool['description'].strip()}" + + # Add parameters section + tool_desc += '\n' + + # Process parameters if they exist in the expected structure + if ('parameters' in tool and isinstance(tool['parameters'], dict) and 'properties' in tool['parameters'] + and isinstance(tool['parameters']['properties'], dict)): + + for param_name, param_fields in tool['parameters']['properties'].items(): + tool_desc += '\n' + tool_desc += f'\n{param_name}' + + if 'type' in param_fields: + tool_desc += f"\n{str(param_fields['type'])}" + + if 'description' in param_fields: + tool_desc += f"\n{param_fields['description'].strip()}" + + # Add any extra parameter fields + handled_param_keys = ['name', 'type', 'description'] + tool_desc += render_extra_keys(param_fields, handled_param_keys) + + tool_desc += '\n' + # Add any extra parameter section fields + handled_keys = ['type', 'properties'] + if 'parameters' in tool: + tool_desc += render_extra_keys(tool['parameters'], handled_keys) + + tool_desc += '\n' + + # Add any extra function fields + handled_keys = ['type', 'name', 'description', 'parameters'] + tool_desc += render_extra_keys(tool, handled_keys) + + tool_desc += '\n' + + tool_descs.append(tool_desc) + + tool_descs.append(TOOL_DESC_SUFFIX) + tool_descs = '\n'.join(tool_descs) + return tool_descs + + def _format_tool_calls(self, tool_call_messages): + result_parts = [] + for idx, message in enumerate(tool_call_messages): + tool_call = self._parse_tool_call(message['content']) + result_parts.append(f"\n\n") + # Processing parameters (if present) + if 'arguments' in tool_call and tool_call['arguments']: + for args_name, args_value in tool_call['arguments'].items(): + result_parts.append(f'\n') + # Handle different types of parameter values + if isinstance(args_value, (dict, list)): + # For dictionaries or lists, use json formatting + args_value = json.dumps(args_value, ensure_ascii=False) + else: + # For other types, convert to strings + args_value = str(args_value) + result_parts.append(f'{args_value}\n\n') + # Close tags + result_parts.append('\n') + # ref: https://github.com/QwenLM/Qwen3-Coder/blob/0ae30f55e9d6c47ff763c334f99c135ad68915dd/qwencoder-eval/tool_calling_eval/berkeley-function-call-leaderboard/bfcl_eval/model_handler/local_inference/qwen_fc.py#L21 # noqa + if idx != len(tool_call_messages) - 1: + result_parts.append('\n') + return ''.join(result_parts) + + def _get_tool_responses(self, tool_messages): + res_tool = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res_tool.append(f'\n{tool_content}\n\n') + return ''.join(res_tool) + + +class Qwen3_5AgentTemplate(Qwen3CoderAgentTemplate): + + def _add_tool_call_prefix(self, tool_content: str, pre_message=None) -> str: + """Qwen3.5/3.6 jinja inserts \n\n between assistant content and + only when effective content (after think removal) is non-empty.""" + if not pre_message or pre_message.get('role') != 'assistant': + return tool_content + content = pre_message.get('content', '') + if not isinstance(content, str): + return tool_content + # Mirror jinja: content.split('')[-1].lstrip('\n') then content|trim + if '' in content: + effective = content.split('')[-1].lstrip('\n') + else: + effective = content + if effective.strip(): + return '\n\n' + tool_content + return tool_content + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools] + tools_prompt = """# Tools + +You have access to the following functions:\n\n +""" + '\n'.join(tool_descs) + f'\n{TOOL_DESC_SUFFIX}' + if system: + tools_prompt += f'\n\n{system}' + return tools_prompt + + def _get_tool_responses(self, tool_messages): + res_tool = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res_tool.append(f'\n{tool_content}\n') + return '\n'.join(res_tool) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/react.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/react.py new file mode 100644 index 0000000000000000000000000000000000000000..b0a2a4d54379f27e3727e98012ba86c826cb957a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/react.py @@ -0,0 +1,72 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import List, Optional, Union + +from .base import BaseAgentTemplate + + +class ReactEnAgentTemplate(BaseAgentTemplate): + + def _format_standalone_tool_responses(self, tool_messages): + return self._format_react_standalone_tool_responses(tool_messages) + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_names = [] + tool_descs = [] + for tool in tools: + tool_desc = self._parse_tool(tool, 'en') + tool_names.append(tool_desc.name_for_model) + tool_descs.append( + f'{tool_desc.name_for_model}: Call this tool to interact with the {tool_desc.name_for_human} API. ' + f'What is the {tool_desc.name_for_human} API useful for? {tool_desc.description_for_model} ' + f'Parameters: {tool_desc.parameters} {tool_desc.args_format}') + + return """Answer the following questions as best you can. You have access to the following tools: + +""" + '\n\n'.join(tool_descs) + f""" + +Use the following format: + +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take, should be one of [{','.join(tool_names)}] +Action Input: the input to the action +Observation: the result of the action +... (this Thought/Action/Action Input/Observation can be repeated zero or more times) +Thought: I now know the final answer +Final Answer: the final answer to the original input question + +Begin! +""" + + +class ReactZnAgentTemplate(BaseAgentTemplate): + + def _format_standalone_tool_responses(self, tool_messages): + return self._format_react_standalone_tool_responses(tool_messages) + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_names = [] + tool_descs = [] + for tool in tools: + tool_desc = self._parse_tool(tool, 'zh') + tool_names.append(tool_desc.name_for_model) + tool_descs.append(f'{tool_desc.name_for_model}: 调用此工具与 {tool_desc.name_for_human} API 进行交互。' + f'{tool_desc.name_for_human} 有什么用？{tool_desc.description_for_model} ' + f'输入参数：{tool_desc.parameters} {tool_desc.args_format}') + return """尽可能地回答以下问题。你可以使用以下工具: + +""" + '\n\n'.join(tool_descs) + f""" + +请按照以下格式进行: + +Question: 需要你回答的输入问题 +Thought: 你应该总是思考该做什么 +Action: 需要使用的工具，应该是[{','.join(tool_names)}]中的一个 +Action Input: 传入工具的内容 +Observation: 行动的结果 +... (这个Thought/Action/Action Input/Observation可以重复N次) +Thought: 我现在知道最后的答案 +Final Answer: 对原始输入问题的最终答案 + +现在开始！ +""" diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/seed_oss.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/seed_oss.py new file mode 100644 index 0000000000000000000000000000000000000000..d838f49a36dbdda4adbaf7c806a294380e82ac12 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/seed_oss.py @@ -0,0 +1,153 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import re +from typing import List, Optional, Tuple, Union + +from swift.infer_engine import Function +from swift.template import Prompt +from .base import BaseAgentTemplate + + +class SeedAgentTemplate(BaseAgentTemplate): + TOOL_CALL_START = '' + TOOL_CALL_END = '' + FUNCTION_TAG = 'function' + PARAMETER_TAG = 'parameter' + + _PY_TYPE_MAPPING = { + 'string': 'str', + 'number': 'int', + 'integer': 'int', + 'boolean': 'bool', + 'array': 'list', + } + + @staticmethod + def _py_type(t: str) -> str: + return SeedAgentTemplate._PY_TYPE_MAPPING.get(t, 'Any') + + def get_toolcall(self, response: str) -> List[Function]: + res_list = re.findall(rf'{self.TOOL_CALL_START}(.+?){self.TOOL_CALL_END}', response, re.DOTALL) + if not res_list: + return super().get_toolcall(response) + + functions = [] + for res in res_list: + func_name_match = re.search(rf'<{self.FUNCTION_TAG}=([^>]+)>', res) + if not func_name_match: + continue + + func_name = func_name_match.group(1) + param_matches = re.findall(rf'<{self.PARAMETER_TAG}=([^>]+)>(.*?)', res, re.DOTALL) + arguments = {name: value for name, value in param_matches} + functions.append(Function(name=func_name, arguments=arguments)) + + return functions + + def _get_tool_responses(self, tool_messages: List[dict]) -> str: + responses = [f"tool\n{tool_message['content']}" for tool_message in tool_messages] + return ''.join(responses) + 'assistant\n' + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages: List[dict], + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + + formatted_tool_responses = self._get_tool_responses(tool_messages) + return assistant_content, ['', formatted_tool_responses] + + def _build_tool_def_string(self, tool: dict) -> str: + """Helper to build a single tool definition string.""" + func = tool.get('function', {}) + func_name = func.get('name') + + if not func_name: + return '' + + parameters = func.get('parameters', {}) + properties = parameters.get('properties', {}) + params = [ + f"{name}: {self._py_type(spec.get('type', 'any'))}" for name, spec in properties.items() + if isinstance(spec, dict) + ] + param_str = ','.join(params) + + docstring_parts = [' """', f' {func.get("description", "").strip()}'] + + if properties: + docstring_parts.append('\n Args:') + required_params = parameters.get('required', []) + for name, spec in properties.items(): + if isinstance(spec, dict): + req_tag = '[必填]' if name in required_params else '[选填]' + desc = spec.get('description', '') + type_str = self._py_type(spec.get('type', 'any')) + docstring_parts.append(f' - {name} ({type_str}) {req_tag}: {desc}') + + returns_props = func.get('returns', {}).get('properties', {}) + if returns_props: + docstring_parts.append('\n Returns:') + for name, spec in returns_props.items(): + desc = spec.get('description', '') + type_str = self._py_type(spec.get('type', 'any')) + docstring_parts.append(f' - {name} ({type_str}): {desc}') + + docstring_parts.append('\n """') + docstring = '\n'.join(docstring_parts) + + return f'Function:\ndef {func_name}({param_str}):\n{docstring}' + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + if not tools: + return system or '' + + tool_defs = [ + tool_def for tool in tools if (wrapped_tool := self.wrap_tool(tool)).get('type') == 'function' and ( + tool_def := self._build_tool_def_string(wrapped_tool)) != '' + ] + tool_defs_joined = '\n\n'.join(tool_defs) + + tool_call_format_instruction = ( + '工具调用请遵循如下格式:\n' + f'{self.TOOL_CALL_START}\n' + f'<{self.FUNCTION_TAG}=example_function_name>\n' + f'<{self.PARAMETER_TAG}=example_parameter_1>value_1\n' + f'<{self.PARAMETER_TAG}=example_parameter_2>This is the value for the second parameter\n' + 'that can span\n' + f'multiple lines\n' + f'\n' + f'{self.TOOL_CALL_END}') + + split_token = 'system' + + if system and split_token in system: + parts = system.split(split_token, 1) + return f'{parts[0]}\n\n{tool_defs_joined}\n{tool_call_format_instruction}\n{split_token}{parts[1]}' + else: + doubao_prompt = ('You are Doubao, a helpful AI assistant. ' + 'You may call one or more functions to assist with the user query.') + return (f'{doubao_prompt}\n\n{tool_defs_joined}\n{tool_call_format_instruction}\n' + f'{split_token}\n{system or ""}') + + def _format_tool_calls(self, tool_call_messages: List[dict]) -> str: + formatted_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + func_name = tool_call['name'] + arguments = tool_call.get('arguments', {}) + + call_parts = [f'<{self.FUNCTION_TAG}={func_name}>'] + for arg_name, arg_value in arguments.items(): + arg_value_str = arg_value if isinstance(arg_value, str) else json.dumps(arg_value, ensure_ascii=False) + call_parts.append(f'<{self.PARAMETER_TAG}={arg_name}>{arg_value_str}') + + call_parts.append(f'') + call_parts_joined = '\n'.join(call_parts) + + full_call = f'{self.TOOL_CALL_START}\n{call_parts_joined}\n{self.TOOL_CALL_END}' + formatted_calls.append(full_call) + return '\n'.join(formatted_calls) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/toolbench.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/toolbench.py new file mode 100644 index 0000000000000000000000000000000000000000..56116e697a524c170aa432730197279e0171296e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/toolbench.py @@ -0,0 +1,38 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +from typing import List, Optional, Union + +from .base import BaseAgentTemplate + + +class ToolBenchAgentTemplate(BaseAgentTemplate): + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + for i, tool in enumerate(tools): + tools[i] = self.unwrap_tool(tool) + tools = json.dumps(tools, ensure_ascii=False) + return f"""You can use many tools(functions) to do the following task. +First I will give you the task description, and your task start. +At each step, you need to give your thought to analyze the status now and what to do next, \ +with a function call to actually execute your step. Your output should follow this format: +Thought: +Action: +Action Input: + +After the call, you will get the call result, and you are now in a new state. +Then you will analyze your status now, then decide what to do next... +After many (Thought-call) pairs, you finally perform the task, then you can give your final answer. +Remember: +1.the state change is irreversible, you can't go back to one of the former state, if you want to restart the task, \ +say \"I give up and restart\". +2.All the thought is short, at most in 5 sentence. +3.You can do more then one try, so if your plan is to continuously try some conditions, \ +you can do one of the conditions per try. +Let's Begin! +Task description: You should use functions to help handle the real time user queries. Remember: +1.ALWAYS call \"Finish\" function at the end of the task. And the final answer should contain enough information \ +to show to the user,If you can't handle the task, \ +or you find that function calls always fail(the function is not valid now), \ +use function Finish->give_up_and_restart. +2.Do not use origin tool names, use only subfunctions' names. +Specifically, you have access to the following APIs: {tools}""" diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/youtu.py b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/youtu.py new file mode 100644 index 0000000000000000000000000000000000000000..2465bb20171aef6ab7dea8e5381dfb0090829781 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/agent_template/youtu.py @@ -0,0 +1,66 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +from typing import TYPE_CHECKING, List, Optional, Tuple, Union + +from .hermes import HermesAgentTemplate + +if TYPE_CHECKING: + from swift.template import Prompt + + +class YoutuAgentTemplate(HermesAgentTemplate): + """Agent template for Youtu-LLM models. + + Tool calling format: + - Tool call: {"name": "function-name", "arguments": {...}} + - Tool response: ... + """ + + def _get_tool_responses(self, tool_messages): + res_tool = [] + for tool_message in tool_messages: + tool_content = tool_message['content'] + res_tool.append(f'{tool_content}') + return '\n'.join(res_tool) + + def _format_tool_responses( + self, + assistant_content: str, + tool_messages, + ) -> Tuple[str, 'Prompt']: + with_action = self.keyword.action in assistant_content and self.keyword.action_input in assistant_content + if with_action: + return super()._format_tool_responses(assistant_content, tool_messages) + # For Youtu-LLM, tool responses are placed in user message + if hasattr(self, 'template_meta'): + prompt = self.template_meta.prompt + chat_sep = self.template_meta.chat_sep + else: + prompt = ['<|User|>{{QUERY}}<|Assistant|>'] + chat_sep = ['<|end_of_text|>'] + res = chat_sep.copy() + total_tool = self._get_tool_responses(tool_messages) + for context in prompt: + if isinstance(context, str): + context = context.replace('{{QUERY}}', total_tool) + res.append(context) + return assistant_content, res + + def _format_tools(self, tools: List[Union[str, dict]], system: Optional[str] = None, user_message=None) -> str: + tool_descs = [json.dumps(self.wrap_tool(tool), ensure_ascii=False) for tool in tools] + system = system or '' + if system: + system = f'{system}\n\n' + return f"""{system}<|begin_of_tool_description|>Tool calling capabilities. +You may call one or more functions to assist with the user query. You have the following functions available: +""" + '\n'.join([f'```json\n{desc}\n```' for desc in tool_descs]) + """ +For tool call returns, you MUST use the following format: +{"name": "function-name", "arguments": {"param1": "value1", "param2": "value2"}} +<|end_of_tool_description|>""" + + def _format_tool_calls(self, tool_call_messages): + tool_calls = [] + for message in tool_call_messages: + tool_call = self._parse_tool_call(message['content']) + tool_calls.append(f'{json.dumps(tool_call, ensure_ascii=False)}') + return ''.join(tool_calls) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..515b69d8edcc9f470b39d27528a488a4987e0a63 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .app_args import AppArguments +from .base_args import BaseArguments, DataArguments, ModelArguments, TemplateArguments, get_supported_tuners +from .deploy_args import DeployArguments, RolloutArguments +from .eval_args import EvalArguments +from .export_args import ExportArguments +from .infer_args import InferArguments +from .pretrain_args import PretrainArguments +from .rlhf_args import RLHFArguments +from .sampling_args import SamplingArguments +from .sft_args import SftArguments +from .tuner_args import TunerArguments +from .webui_args import WebUIArguments diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/app_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/app_args.py new file mode 100644 index 0000000000000000000000000000000000000000..00b63181a2240ec6dea73b8c6775cce488f7936c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/app_args.py @@ -0,0 +1,57 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import dataclass +from typing import Literal, Optional + +from swift.model import get_matched_model_meta +from swift.template import get_template_meta +from swift.utils import find_free_port, get_logger +from .deploy_args import DeployArguments +from .webui_args import WebUIArguments + +logger = get_logger() + + +@dataclass +class AppArguments(WebUIArguments, DeployArguments): + """Arguments for configuring the Web UI inference. + + This dataclass inherits from WebUIArguments and DeployArguments, combining their settings to configure the user + interface for model inference. + + Args: + base_url (Optional[str]): The base URL for the model deployment API, e.g., `http://localhost:8000/v1`. If set + to `None`, a local deployment will be used instead. Defaults to None. + studio_title (Optional[str]): The title for the Web UI studio. If set to `None`, the title will default to the + model's name. Defaults to None. + is_multimodal (Optional[bool]): Whether to launch the multimodal version of the application. If `None`, the + app will attempt to auto-detect this setting based on the model. If auto-detection is not possible, it + defaults to `False`. Defaults to None. + lang (str): Overrides the language setting for the Web UI. Defaults to 'en'. + verbose (bool): Whether to log detailed request information. Defaults to False. + stream (bool): Whether to enable streaming output for model responses. Defaults to True. + """ + base_url: Optional[str] = None + studio_title: Optional[str] = None + is_multimodal: Optional[bool] = None + + lang: Literal['en', 'zh'] = 'en' + verbose: bool = False + stream: bool = True + + def _init_torch_dtype(self) -> None: + if self.base_url: + self.model_meta = get_matched_model_meta(self.model) + self.model_info = None + return + super()._init_torch_dtype() + + def __post_init__(self): + DeployArguments.__post_init__(self) + self.server_port = find_free_port(self.server_port) + if self.model_meta: + if self.system is None: + self.system = get_template_meta(self.model_info, self.model_meta).default_system + if self.is_multimodal is None: + self.is_multimodal = self.model_meta.is_multimodal + if self.is_multimodal is None: + self.is_multimodal = False diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eaf3f0d03c266a78c2ee1f24e83fd359893e1a7d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base_args import BaseArguments, get_supported_tuners +from .data_args import DataArguments +from .generation_args import GenerationArguments +from .model_args import ModelArguments +from .quant_args import QuantizeArguments +from .template_args import TemplateArguments diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/base_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/base_args.py new file mode 100644 index 0000000000000000000000000000000000000000..326fb6d8b7824461d161eacb097f24daaa70d458 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/base_args.py @@ -0,0 +1,367 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import os +import peft +import shutil +from dataclasses import dataclass, field, fields +from packaging import version +from typing import Any, Dict, List, Literal, Optional, Union + +import swift +from swift.dataset import load_dataset +from swift.hub import get_hub +from swift.model import get_ckpt_dir, get_model_processor, load_by_unsloth +from swift.ray_utils import RayArguments +from swift.template import Template, get_template +from swift.tuner_plugin import tuners_map +from swift.utils import (Processor, check_json_format, get_dist_setting, get_logger, import_external_file, is_dist, + is_master, json_parse_to_dict, safe_snapshot_download, set_device, use_hf_hub) +from .data_args import DataArguments +from .generation_args import GenerationArguments +from .model_args import ModelArguments +from .quant_args import QuantizeArguments +from .template_args import TemplateArguments + +logger = get_logger() + + +def get_supported_tuners(): + return {'lora', 'full', 'longlora', 'adalora', 'llamapro', 'adapter', 'vera', 'boft', 'fourierft', 'reft', 'bone' + } | set(tuners_map.keys()) + + +def _patch_peft(): + """Patch peft functions that are incompatible with SWIFT. + + 1. _maybe_shard_state_dict_for_tp: TP sharding is not used by SWIFT, and causes errors + when torch.distributed is initialized (e.g. MoE training with target_parameters). + 2. _maybe_shard_state_dict_for_tp internal logic accesses base_layer.weight.device which + fails for expert modules that don't have a `weight` attribute. + """ + if version.parse(peft.__version__) >= version.parse('0.19.0'): + from peft.utils import save_and_load + save_and_load._maybe_shard_state_dict_for_tp = lambda model, state_dict, adapter_name: None + + +@dataclass +class BaseArguments(GenerationArguments, QuantizeArguments, DataArguments, TemplateArguments, ModelArguments, + RayArguments): + """BaseArguments class is a dataclass that inherits from multiple argument classes. + + This class consolidates arguments from GenerationArguments, QuantizeArguments, DataArguments, + TemplateArguments, ModelArguments, RayArguments. + + Args: + tuner_backend (str): The tuner backend to use. Choices are 'peft' or 'unsloth'. Default is 'peft'. + tuner_type (str): The tuner type. Choices include 'lora', 'full', 'longlora', 'adalora', 'llamapro', + 'adapter', 'vera', 'boft', 'fourierft', 'reft'. Default is 'lora'. + adapters (List[str]): A list of adapter IDs or paths. This is typically used for inference or deployment. + It can also resume training by only loading adapter weights, differing from `resume_from_checkpoint` + which also loads optimizer states. Default is []. + external_plugins (List[str]): A list of external 'plugin.py' files to be registered and imported into + the plugin module. Default is []. + seed (int): The global random seed for reproducibility. Note that this does not affect `data_seed`, + which controls dataset randomization. Default is 42. + model_kwargs (Optional[str]): Additional keyword arguments for specific models, passed as a JSON string + (e.g., '{"key": "value"}'). It's recommended to use the same arguments for inference as for training. + Default is None. + enable_npu_model_patch (bool): Whether to enable model-related NPU patches. Default is True. + load_args (bool): Whether to load `args.json` from a checkpoint when using `--resume_from_checkpoint`, + `--model`, or `--adapters`. Defaults to True for inference/export and False for training. Usually, + this does not need to be modified. Default is True. + load_data_args (bool): If True, will also load data-related arguments from `args.json`. This is useful + for running inference on the same validation split used during training. Default is False. + packing (bool): Whether to enable packing of datasets. Default is False. + packing_length (Optional[int]): Length of packing. Default is None. + packing_num_proc (int): Number of processes used for packing, Default is 1. + packing_strategy (Literal['binpack', 'sequential']): Packing algorithm. 'binpack' (default) uses + best-fit-decreasing bin packing (reorders samples); 'sequential' uses order-preserving greedy + packing (next-fit: a single open pack, flushed when the next sample doesn't fit) so the sample + order / pack boundaries follow a sequential sampler (use packing_num_proc=1). Default is 'binpack'. + lazy_tokenize (Optional[bool]): Whether to enable lazy tokenization. Default is None. + use_hf (bool): Whether to use Hugging Face for downloading/uploading models and datasets. If False, + ModelScope is used. Default is False. + hub_token (Optional[str]): The authentication token for ModelScope or Hugging Face Hub. Default is None. + ddp_timeout (int): Timeout for DDP (Distributed Data Parallel) operations, in seconds. Default is 18000000. + ddp_backend (Optional[str]): The backend for DDP. Choices include "nccl", "gloo", "mpi", "ccl", "hccl", + "cncl", "mccl". If None, it will be automatically selected. Default is None. + ignore_args_error (bool): Whether to ignore argument errors. This is useful for compatibility with Jupyter + notebooks. Default is False. + use_swift_lora (bool): Whether to use swift lora. This is a compatible argument. Default is False. + """ + tuner_backend: Literal['peft', 'unsloth'] = 'peft' + tuner_type: str = field(default='lora', metadata={'help': f'tuner_type choices: {list(get_supported_tuners())}'}) + adapters: List[str] = field(default_factory=list) + external_plugins: List[str] = field(default_factory=list) + # This parameter is kept for swift3.x compatibility. Please use `external_plugins` as a replacement. + custom_register_path: List[str] = field(default_factory=list) + + seed: int = 42 + model_kwargs: Optional[Union[dict, str]] = None + enable_npu_model_patch: bool = True + load_args: bool = True + load_data_args: bool = False + # dataset + packing: bool = False + packing_length: Optional[int] = None + packing_num_proc: int = 1 + packing_strategy: Literal['binpack', 'sequential'] = 'binpack' + lazy_tokenize: Optional[bool] = None + # hub + use_hf: bool = False + # None: use env var `MODELSCOPE_API_TOKEN` + hub_token: Optional[str] = field( + default=None, metadata={'help': 'SDK token can be found in https://modelscope.cn/my/myaccesstoken'}) + # dist + ddp_timeout: int = 18000000 + ddp_backend: Optional[str] = None + + # extra + ignore_args_error: bool = False # True: notebook compatibility + use_swift_lora: bool = False # True for using tuner_backend == swift, don't specify this unless you know what you are doing # noqa + + def _prepare_training_args(self, training_args: Dict[str, Any]) -> None: + pass + + def _init_lazy_tokenize(self): + if self.lazy_tokenize is None: + if self.cached_dataset or self.cached_val_dataset: + self.lazy_tokenize = False + elif (self.model_meta is not None and self.model_meta.is_multimodal and not self.streaming + and not self.packing and not getattr(self, 'group_by_length', False)): + self.lazy_tokenize = True + else: + self.lazy_tokenize = False + logger.info(f'Setting args.lazy_tokenize: {self.lazy_tokenize}') + if self.lazy_tokenize: + if self.packing: + raise ValueError('Packing and lazy_tokenize are incompatible.') + if self.streaming: + raise ValueError('Streaming and lazy_tokenize are incompatible.') + + def _import_external_plugins(self): + if isinstance(self.external_plugins, str): + self.external_plugins = [self.external_plugins] + # swift v3.x compatibility + if isinstance(self.custom_register_path, str): + self.custom_register_path = [self.custom_register_path] + if self.custom_register_path: + self.external_plugins += self.custom_register_path + + if not self.external_plugins: + return + for external_plugin in self.external_plugins: + import_external_file(external_plugin) + logger.info(f'Successfully imported external_plugins: {self.external_plugins}.') + + @staticmethod + def _check_is_adapter(adapter_dir: str) -> bool: + if (os.path.exists(os.path.join(adapter_dir, 'adapter_config.json')) + or os.path.exists(os.path.join(adapter_dir, 'default', 'adapter_config.json')) + or os.path.exists(os.path.join(adapter_dir, 'reft'))): + return True + return False + + def _init_adapters(self): + if isinstance(self.adapters, str): + self.adapters = [self.adapters] + self.adapters = [ + safe_snapshot_download(adapter, use_hf=self.use_hf, hub_token=self.hub_token) for adapter in self.adapters + ] + + def __post_init__(self): + _patch_peft() + self.swift_version = swift.__version__ + if self.use_hf or use_hf_hub(): + self.use_hf = True + os.environ['USE_HF'] = '1' + self._init_adapters() + self._init_ckpt_dir() + self._import_external_plugins() + self._init_model_kwargs() + # The Seq2SeqTrainingArguments has a property called world_size, which cannot be assigned a value. + self.rank, self.local_rank, self.global_world_size, self.local_world_size = get_dist_setting() + logger.info(f'rank: {self.rank}, local_rank: {self.local_rank}, ' + f'world_size: {self.global_world_size}, local_world_size: {self.local_world_size}') + if self.tuner_type not in tuners_map: # build-in tuner + for adapter in self.adapters: + assert self._check_is_adapter(adapter), ( + f'`{adapter}` is not an adapter, please try using `--model` to pass it.') + ModelArguments.__post_init__(self) + QuantizeArguments.__post_init__(self) + TemplateArguments.__post_init__(self) + DataArguments.__post_init__(self) + RayArguments.__post_init__(self) + self._init_stream() + if self.max_length is None and self.model_info is not None: + self.max_length = self.model_info.max_model_len + if self.packing and self.packing_length is None: + self.packing_length = self.max_length + self._init_lazy_tokenize() + self.hub = get_hub(self.use_hf) + if self.hub.try_login(self.hub_token): + logger.info('hub login successful!') + + def _init_model_kwargs(self): + """Prepare model kwargs and set them to the env""" + self.model_kwargs: Dict[str, Any] = json_parse_to_dict(self.model_kwargs) + for k, v in self.model_kwargs.items(): + k = k.upper() + os.environ[k] = str(v) + + @property + def is_adapter(self) -> bool: + return self.tuner_type not in {'full'} + + @property + def supported_tuners(self): + return get_supported_tuners() + + @property + def adapters_can_be_merged(self): + return {'lora', 'longlora', 'llamapro', 'adalora'} + + @classmethod + def from_pretrained(cls, checkpoint_dir: str): + self = super().__new__(cls) + self.load_data_args = True + self.ckpt_dir = checkpoint_dir + self.load_args_from_ckpt() + all_keys = list(f.name for f in fields(BaseArguments)) + for key in all_keys: + if not hasattr(self, key): + setattr(self, key, None) + return self + + def _init_ckpt_dir(self, adapters=None): + # compat megatron + model = self.model or getattr(self, 'mcore_model', None) + adapters = adapters or self.adapters or getattr(self, 'mcore_adapter', None) + if isinstance(adapters, str): + adapters = [adapters] + self.ckpt_dir = get_ckpt_dir(model, adapters) + if self.ckpt_dir and self.load_args: + self.load_args_from_ckpt() + + def load_args_from_ckpt(self) -> None: + args_path = os.path.join(self.ckpt_dir, 'args.json') + assert os.path.exists(args_path), f'args_path: {args_path}' + with open(args_path, 'r', encoding='utf-8') as f: + old_args = json.load(f) + force_load_keys = [ + # base_args + 'tuner_type', + # model_args + 'task_type', + # quant_args + 'bnb_4bit_quant_type', + 'bnb_4bit_use_double_quant', + ] + # If the current value is None or an empty list and it is among the following keys + load_keys = [ + 'external_plugins', + # model_args + 'model', + 'model_type', + 'model_revision', + 'torch_dtype', + 'attn_impl', + 'experts_impl', + 'new_special_tokens', + 'num_labels', + 'problem_type', + 'rope_scaling', + 'max_model_len', + # quant_args + 'quant_method', + 'quant_bits', + 'hqq_axis', + 'bnb_4bit_compute_dtype', + # template_args + 'template', + 'system', + 'truncation_strategy', + 'agent_template', + 'norm_bbox', + 'use_chat_template', + 'response_prefix', + ] + data_keys = list(f.name for f in fields(DataArguments)) + swift_version = old_args.get('swift_version') + if swift_version is None or version.parse(swift_version) < version.parse('4.0.0.dev'): + load_keys.remove('model_type') + for key, old_value in old_args.items(): + if old_value is None: + continue + if key in force_load_keys or self.load_data_args and key in data_keys: + setattr(self, key, old_value) + value = getattr(self, key, None) + if key in load_keys and (value is None or isinstance(value, (list, tuple)) and len(value) == 0): + setattr(self, key, old_value) + logger.info(f'Successfully loaded {args_path}.') + + def save_args(self, output_dir=None) -> None: + if is_master(): + output_dir = output_dir or self.output_dir + os.makedirs(output_dir, exist_ok=True) + fpath = os.path.join(output_dir, 'args.json') + logger.info(f'The {self.__class__.__name__} will be saved in: {fpath}') + with open(fpath, 'w', encoding='utf-8') as f: + json.dump(check_json_format(self.__dict__), f, ensure_ascii=False, indent=2) + config_file = os.getenv('SWIFT_CONFIG_FILE') + if config_file: + shutil.copy(config_file, output_dir) + + def _init_device(self): + if is_dist(): + set_device() + + def get_template(self, processor: Optional[Processor] = None, **kwargs) -> Template: + if processor is None: + processor = self.get_model_processor(load_model=False)[1] + template_kwargs = self.get_template_kwargs() + if 'template_type' in kwargs: + template_type = kwargs.get('template_type') + else: + template_type = self.template + template_kwargs['template_type'] = template_type + template = get_template(processor, **template_kwargs) + return template + + def get_model_processor(self, + *, + model=None, + model_type=None, + revision=None, + task_type=None, + num_labels=None, + **kwargs): + if self.tuner_backend == 'unsloth': + return load_by_unsloth(self) + res = self.get_model_kwargs() + res.update(kwargs) + # compat rlhf + res['model_id_or_path'] = model or self.model + res['model_type'] = model_type or self.model_type + res['revision'] = revision or self.model_revision + res['task_type'] = task_type or self.task_type + res['num_labels'] = num_labels or self.num_labels + + return get_model_processor(**res) + + def load_dataset(self): + dataset_kwargs = self.get_dataset_kwargs() + train_dataset, val_dataset = None, None + if self.dataset: + train_dataset, val_dataset = load_dataset( + self.dataset, + split_dataset_ratio=self.split_dataset_ratio, + shuffle=self.dataset_shuffle, + **dataset_kwargs) + if len(self.val_dataset) > 0: + # Loading val dataset + dataset_kwargs.pop('interleave_prob', None) + _, val_dataset = load_dataset( + self.val_dataset, split_dataset_ratio=1.0, shuffle=self.val_dataset_shuffle, **dataset_kwargs) + assert self.split_dataset_ratio == 0. + return train_dataset, val_dataset diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/data_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/data_args.py new file mode 100644 index 0000000000000000000000000000000000000000..01b8b0f7a45c533e54e3aabca0c9a3e255760d58 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/data_args.py @@ -0,0 +1,145 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import dataclass, field +from typing import List, Literal, Optional, Union + +from swift.dataset import register_dataset_info +from swift.utils import get_logger, json_parse_to_dict + +logger = get_logger() + + +@dataclass +class DataArguments: + """Holds arguments related to dataset handling and processing. + + Args: + dataset (List[str]): A list of dataset IDs or paths. Defaults to []. + Format for each dataset: 'dataset_id_or_path:subset#count'. Both subset and count are optional. + - Subsets: Only effective for dataset IDs or folders. Use '/' to select multiple subsets (e.g., + 'dataset_id:subset1/subset2') or 'all' to select all registered subsets. If only one subset is + registered, it will be used by default; otherwise, 'default' is the default. + - Sampling Count: By default, the full dataset is used. Use '#count' to sample. If count < + total samples, it performs random sampling without replacement. If count > total, it repeats + the full dataset `count // total` times and then randomly samples an additional `count % total` + samples. Note: Streaming datasets or setting `--dataset_shuffle false` will result in sequential + sampling. + - Local datasets: Supports formats like jsonl, csv, json, and folders. + val_dataset (List[str]): A list of validation dataset IDs or paths. Defaults to []. + cached_dataset (List[str]): Use cached datasets to avoid GPU time being occupied by tokenization during + training/inference on large datasets. This parameter is used to set the folder path(s) of + cached training datasets, and defaults to `[]`. + This is generated by the `swift export --to_cached_dataset true ...` command. + ms-swift only stores an extra 'length' field and filters out erroneous samples + to reduce storage. Actual preprocessing happens concurrently with training. + cached_val_dataset (List[str]): Folder path(s) for cached validation datasets, default is []. + split_dataset_ratio (float): The ratio to split from the training set for validation if `val_dataset` is not + provided. Defaults to 0.0. Note: The default was 0.01 in `ms-swift<3.6`. + data_seed (int): The random seed for dataset shuffling. Defaults to 42. + dataset_num_proc (int): The number of processes to use for dataset preprocessing. Defaults to 1. + load_from_cache_file (bool): Whether to load the dataset from cache files. Recommended to set to `True` during + actual runs and `False` during debugging. Defaults to False. + Note: The default was `True` in `ms-swift<3.9`. + dataset_shuffle (bool): Whether to shuffle the training dataset. Defaults to True. + Note: For CPT/SFT, shuffling occurs at both the dataset level (controlled by this flag) and the dataloader + level. + val_dataset_shuffle (bool): Whether to shuffle the validation dataset. Defaults to False. + streaming (bool): Enables streaming to read and process the dataset on-the-fly. `--max_steps` must be set as the + dataset length is unknown. This allows preprocessing to overlap with training but can become a bottleneck + with a large `world_size` as preprocessing only runs on rank 0. Defaults to False. + interleave_prob (Optional[List[float]]): If set, combines datasets using `interleave_datasets` with the + provided probabilities instead of `concatenate_datasets`. Typically used for streaming. Defaults to None. + stopping_strategy (str): The stopping strategy for `interleave_datasets`. Can be "first_exhausted" or + "all_exhausted". Defaults to "first_exhausted". + shuffle_buffer_size (int): The buffer size for shuffling in streaming mode. Only effective if `dataset_shuffle` + is `True`. Defaults to 1000. + download_mode (str): The dataset download mode. Options are 'reuse_dataset_if_exists' and 'force_redownload'. + Defaults to 'reuse_dataset_if_exists'. + columns (Optional[str]): A JSON string for column mapping to fit the format required by `AutoPreprocessor`. + Example: '{"text1": "query", "text2": "response"}'. Defaults to None. + strict (bool): If `True`, raises an error on any problematic data row. If `False`, discards the problematic + sample and continues. Typically used for debugging. Defaults to False. + remove_unused_columns (bool): Whether to remove columns not used by the model. If `False`, extra columns are + passed to the trainer's `compute_loss` function, which is useful for custom loss calculations. + Defaults to True. Note: The default is `False` for GPRO. + disable_auto_column_mapping (bool): By default, column names in the dataset are automatically mapped. + This parameter disables that behavior (the `columns` parameter remains effective), defaulting to `False`. + model_name (Optional[List[str]]): For self-cognition tasks, replaces the `{{NAME}}` placeholder in the + `swift/self-cognition` dataset. Pass Chinese and English names. + Example: `--model_name 小黄 'Xiao Huang'`. Defaults to None. + model_author (Optional[List[str]]): For self-cognition tasks, replaces the `{{AUTHOR}}` placeholder in the + `swift/self-cognition` dataset. Pass author's Chinese and English names. + Example: `--model_author '魔搭' 'ModelScope'`. Defaults to None. + custom_dataset_info (List[str]): Path to a custom dataset registration JSON file. Defaults to []. + """ + # dataset_id or dataset_dir or dataset_path + dataset: List[str] = field(default_factory=list) + val_dataset: List[str] = field(default_factory=list) + cached_dataset: List[str] = field(default_factory=list) + cached_val_dataset: List[str] = field(default_factory=list) + split_dataset_ratio: float = 0. + + data_seed: int = 42 + dataset_num_proc: int = 1 + load_from_cache_file: bool = False + dataset_shuffle: bool = True + val_dataset_shuffle: bool = False + streaming: bool = False + interleave_prob: Optional[List[float]] = None + stopping_strategy: Literal['first_exhausted', 'all_exhausted'] = 'first_exhausted' + shuffle_buffer_size: int = 1000 + + download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists' + columns: Optional[Union[dict, str]] = None + strict: bool = False + remove_unused_columns: bool = True + disable_auto_column_mapping: bool = False + # Chinese name and English name + model_name: Optional[List[str]] = field(default=None, metadata={'help': "e.g. ['小黄', 'Xiao Huang']"}) + model_author: Optional[List[str]] = field(default=None, metadata={'help': "e.g. ['魔搭', 'ModelScope']"}) + + custom_dataset_info: List[str] = field(default_factory=list) # .json + + def _init_custom_dataset_info(self): + """register custom dataset_info.json to datasets""" + if isinstance(self.custom_dataset_info, str): + self.custom_dataset_info = [self.custom_dataset_info] + for path in self.custom_dataset_info: + register_dataset_info(path) + + def __post_init__(self): + self.columns = json_parse_to_dict(self.columns) + if len(self.val_dataset) > 0 or self.streaming and self.split_dataset_ratio > 0: + self.split_dataset_ratio = 0. + if len(self.val_dataset) > 0: + msg = 'len(args.val_dataset) > 0' + else: + msg = 'args.streaming is True' + logger.info(f'Because {msg}, setting split_dataset_ratio: {self.split_dataset_ratio}') + self._init_custom_dataset_info() + if isinstance(self.cached_dataset, str): + self.cached_dataset = [self.cached_dataset] + self._init_val_dataset_exists() + + def _init_val_dataset_exists(self): + self._val_dataset_exists = bool(self.dataset and self.split_dataset_ratio > 0 or self.val_dataset + or self.cached_val_dataset) + + def get_dataset_kwargs(self): + return { + 'seed': self.data_seed, + 'num_proc': self.dataset_num_proc, + 'load_from_cache_file': self.load_from_cache_file, + 'streaming': self.streaming, + 'interleave_prob': self.interleave_prob, + 'stopping_strategy': self.stopping_strategy, + 'shuffle_buffer_size': self.shuffle_buffer_size, + 'use_hf': self.use_hf, + 'hub_token': self.hub_token, + 'download_mode': self.download_mode, + 'columns': self.columns, + 'strict': self.strict, + 'model_name': self.model_name, + 'model_author': self.model_author, + 'remove_unused_columns': self.remove_unused_columns, + 'disable_auto_column_mapping': self.disable_auto_column_mapping, + } diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/generation_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/generation_args.py new file mode 100644 index 0000000000000000000000000000000000000000..5dd535e2af457848916c70323176cc1ee0c1db1a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/generation_args.py @@ -0,0 +1,76 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import dataclass, field +from typing import List, Optional + +from swift.infer_engine import RequestConfig +from swift.utils import get_logger + +logger = get_logger() + + +@dataclass +class GenerationArguments: + """A dataclass that holds arguments for text generation. + + Args: + max_new_tokens (Optional[int]): The maximum number of new tokens to generate. Defaults to None (unlimited). + temperature (Optional[float]): The sampling temperature. A higher temperature makes the output more random. To + disable randomness, you can set this to 0 or `top_k` to 1. Defaults to None, which means loading from + 'generation_config.json'. + top_k (Optional[int]): The number of highest probability vocabulary tokens to keep for top-k-filtering. + Defaults to None (reads from 'generation_config.json'). + top_p (Optional[float]): The cumulative probability for nucleus sampling. Filters the vocabulary to the + smallest set of tokens whose cumulative probability exceeds `top_p`. Defaults to None (reads from + 'generation_config.json'). + repetition_penalty (Optional[float]): The penalty applied to repeated tokens. A value of 1.0 means no penalty. + Defaults to None (reads from 'generation_config.json'). + num_beams (Optional[int]): The number of beams to use for beam search. Defaults to 1. + stream (bool): Whether to enable streaming output. Defaults to None, which is `True` for interactive mode and + `False` for batch inference. Note: For ms-swift < 3.6, the default is `False`. + stop_words (List[str]): A list of extra stop words, in addition to the end-of-sequence token. Note: The + `eos_token` is removed from the output, while these stop words are preserved. Defaults to an empty list. + logprobs (bool): Whether to output log probabilities of the generated tokens. Defaults to False. + top_logprobs (Optional[int]): The number of top log probabilities to return for each token position. Requires + `logprobs` to be True. Defaults to None. + structured_outputs_regex (Optional[str]): A regular expression pattern for structured outputs (guided decoding). + When set, the model's generation is constrained to match the specified regex pattern. This is useful for + tasks requiring structured outputs like reasoning chains. Only effective when `infer_backend` is 'vllm'. + Defaults to None. + """ + + # generation config + max_new_tokens: Optional[int] = None # Unlimited, constrained by max_model_len. + # If it is None, use the parameters from generation_config. + temperature: Optional[float] = None # Set to 0, which means do_sample is False. + top_k: Optional[int] = None + top_p: Optional[float] = None + repetition_penalty: Optional[float] = None + num_beams: int = 1 + + stream: Optional[bool] = None + stop_words: List[str] = field(default_factory=list) + logprobs: bool = False + top_logprobs: Optional[int] = None + # structured outputs (guided decoding), only effective for vllm backend + structured_outputs_regex: Optional[str] = None + + def _init_stream(self): + if self.stream is None: + self.stream = False + + def get_request_config(self): + if getattr(self, 'task_type') != 'causal_lm': + return + + return RequestConfig( + max_tokens=self.max_new_tokens, + temperature=self.temperature, + top_p=self.top_p, + top_k=self.top_k, + num_beams=self.num_beams, + stop=self.stop_words, + stream=self.stream, + repetition_penalty=self.repetition_penalty, + logprobs=self.logprobs, + top_logprobs=self.top_logprobs, + structured_outputs_regex=self.structured_outputs_regex) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/model_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/model_args.py new file mode 100644 index 0000000000000000000000000000000000000000..4caeed8802c78a6d1f94f156a01fc80eb9c18a3f --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/model_args.py @@ -0,0 +1,249 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import ast +import math +import os +import torch +from dataclasses import dataclass, field +from transformers.utils import is_torch_mps_available +from typing import Any, Dict, List, Literal, Optional, Union + +from swift.model import MODEL_MAPPING, get_model_info_meta, get_model_name +from swift.utils import HfConfigFactory, get_dist_setting, get_logger, json_parse_to_dict + +logger = get_logger() + + +@dataclass +class ModelArguments: + """A dataclass that holds various arguments related to model configuration and usage. + + Args: + model (Optional[str]): The model ID from the Hub or a local path to the model. Defaults to None. + model_type (Optional[str]): The model type. In ms-swift, a 'model_type' groups models with the same + architecture, loading process, and template. Defaults to None, which enables auto-selection based on + the suffix of `--model` and the 'architectures' attribute in `config.json`. The `model_type` for a + corresponding model can be found in the list of supported models. Note: The concept of `model_type` + in ms-swift differs from the `model_type` in `config.json`. Custom models usually require registering + their own `model_type` and `template`. + model_revision (Optional[str]): The revision of the model. Defaults to None. + task_type (str): The task type. Can be 'causal_lm', 'seq_cls', 'embedding', 'reranker', or + 'generative_reranker'. If set to 'seq_cls', you usually need to specify `--num_labels` and + `--problem_type`. Defaults to 'causal_lm'. + torch_dtype (Optional[str]): The data type of the model weights. Supports 'float16', 'bfloat16', 'float32'. + Defaults to None, in which case it's read from the 'config.json' file. + attn_impl (Optional[str]): The attention implementation to use. Options include 'sdpa', 'eager', 'flash_attn', + 'flash_attention_2', 'flash_attention_3', 'flash_attention_4', etc. + Defaults to None, which means it will be read from 'config.json'. + Note: Support for these implementations depends on the model's transformers implementation. + If set to 'flash_attn' (for backward compatibility), 'flash_attention_2' will be used. + experts_impl (Optional[str]): Expert implementation type, options are 'grouped_mm', 'batched_mm', 'eager'. + Defaults to None. This feature requires "transformers>=5.0.0". + new_special_tokens (List[str]): Additional special tokens to be added to the tokenizer. Can also be a path to + a `.txt` file, where each line is a special token. Defaults to an empty list `[]`. + num_labels (Optional[int]): The number of labels for classification tasks (when `--task_type` is 'seq_cls'). + Required for such tasks. Defaults to None. + problem_type (Optional[str]): The problem type for classification tasks (`--task_type` 'seq_cls'). Options are + 'regression', 'single_label_classification', 'multi_label_classification'. Defaults to None, but is + automatically set to 'regression' if the model is a reward_model or `num_labels` is 1, and + 'single_label_classification' otherwise. + rope_scaling (Optional[str]): The RoPE scaling type. You can pass a string like 'linear', 'dynamic', or + 'yarn', and ms-swift will automatically set the corresponding `rope_scaling` and override the + 'config.json' value. Alternatively, you can pass a JSON string (e.g., '{"factor":2.0, "type":"yarn"}'), + which will directly override the `rope_scaling` in 'config.json'. Defaults to None. + device_map (Optional[str]): The device map configuration for the model, e.g., 'auto', 'cpu', a JSON string, + or a path to a JSON file. This argument is passed directly to the `from_pretrained` method of transformers. + Defaults to None, and will be set automatically based on the device and distributed training settings. + max_memory (Optional[str]): The maximum memory allocation for each device when `device_map` is 'auto' or + 'sequential'. Example: '{0: "20GB", 1: "20GB"}'. This argument is passed directly to the `from_pretrained` + method of transformers. Defaults to None. + max_model_len (Optional[int]): The maximum model length. This is used to calculate the RoPE scaling factor + when `rope_scaling` is specified as a string. If not None, it overrides the `max_position_embeddings` + value in 'config.json'. Defaults to None. + local_repo_path (Optional[str]): Path to a local repository for models that require a GitHub repo during + loading (e.g., deepseek-vl2). This avoids network issues during `git clone`. Defaults to None. + init_strategy (Optional[str]): The strategy to initialize all uninitialized parameters when loading a model + (especially for custom architectures). Options include 'zero', 'uniform', 'normal', 'xavier_uniform', + 'xavier_normal', 'kaiming_uniform', 'kaiming_normal', 'orthogonal'. Defaults to None. + """ + model: Optional[str] = None # model id or model path + model_type: Optional[str] = field( + default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'}) + model_revision: Optional[str] = None + task_type: Literal['causal_lm', 'seq_cls', 'embedding', 'reranker', 'generative_reranker'] = None + + torch_dtype: Literal['bfloat16', 'float16', 'float32', None] = None + # flash_attn: It will automatically convert names based on the model. + # None: It will be automatically selected between sdpa and eager. + # 'flash_attn', 'sdpa', 'eager', 'flex_attention', + # 'flash_attention_2', 'flash_attention_3', 'flash_attention_4' + attn_impl: Optional[str] = None + experts_impl: Optional[str] = None + new_special_tokens: List[str] = field(default_factory=list) + + num_labels: Optional[int] = None + problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] = None + rope_scaling: Optional[str] = None + device_map: Optional[Union[dict, str]] = None + max_memory: Optional[Union[dict, str]] = None + max_model_len: Optional[int] = None + # When some model code needs to be downloaded from GitHub, + # this parameter specifies the path to the locally downloaded repository. + local_repo_path: Optional[str] = None + init_strategy: Literal['zero', 'uniform', 'normal', 'xavier_uniform', 'xavier_normal', 'kaiming_uniform', + 'kaiming_normal', 'orthogonal'] = None + + def _init_device_map(self): + """Prepare device map args""" + if self.device_map: + self.device_map: Union[str, Dict[str, Any], None] = json_parse_to_dict(self.device_map, strict=False) + # compat mp&ddp + _, local_rank, _, local_world_size = get_dist_setting() + if local_world_size > 1 and isinstance(self.device_map, dict) and local_rank > 0: + for k, v in self.device_map.items(): + if isinstance(v, int): + self.device_map[k] += local_rank + + def _init_max_memory(self): + if isinstance(self.max_memory, str): + try: + self.max_memory = ast.literal_eval(self.max_memory) + except Exception: + pass + self.max_memory = json_parse_to_dict(self.max_memory) + # compat mp&ddp + _, local_rank, _, local_world_size = get_dist_setting() + if local_world_size > 1 and isinstance(self.max_memory, dict) and local_rank > 0: + for k in list(self.max_memory.keys()): + if isinstance(k, int): + self.max_memory[k + local_rank] = self.max_memory.pop(k) + + def _init_torch_dtype(self) -> None: + """"If torch_dtype is None, find a proper dtype by the config.json/GPU""" + from ..sft_args import SftArguments + + self.torch_dtype: Optional[torch.dtype] = HfConfigFactory.to_torch_dtype(self.torch_dtype) + self.torch_dtype: torch.dtype = self._init_model_info() + # Mixed Precision Training + if isinstance(self, SftArguments): + self._init_mixed_precision() + + def _init_mixed_precision(self): + if is_torch_mps_available(): + fp16, bf16 = False, False + elif self.torch_dtype in {torch.float16, torch.float32}: + fp16, bf16 = True, False + elif self.torch_dtype == torch.bfloat16: + fp16, bf16 = False, True + else: + raise ValueError(f'args.torch_dtype: {self.torch_dtype}') + if self.fp16 is None: + self.fp16 = fp16 + if self.bf16 is None: + self.bf16 = bf16 + + def _init_rope_scaling(self): + if self.rope_scaling: + rope_scaling: dict = json_parse_to_dict(self.rope_scaling, strict=False) + if isinstance(rope_scaling, str): + assert rope_scaling in ['linear', 'dynamic', 'yarn'] + rope_scaling = {'type': rope_scaling} + else: + rope_scaling = self.model_info.rope_scaling + # reset the factor + rope_scaling.pop('factor', None) + + rope_type = rope_scaling.get('rope_type', rope_scaling.get('type', 'default')) + if 'factor' not in rope_scaling and self.max_model_len is None and rope_type == 'default': + # fix megatron qwen2_5_vl + self.rope_scaling = rope_scaling + logger.info(f'Setting args.rope_scaling: {rope_scaling}') + return + + # get origin_max_model_len + origin_max_model_len = None + if rope_scaling and rope_scaling.get('original_max_position_embeddings') is not None: + origin_max_model_len = rope_scaling['original_max_position_embeddings'] + elif self.model_info.rope_scaling: + if self.model_info.rope_scaling.get('original_max_position_embeddings') is not None: + origin_max_model_len = self.model_info.rope_scaling['original_max_position_embeddings'] + elif self.model_info.rope_scaling.get('factor') is not None: + origin_max_model_len = self.model_info.max_model_len // self.model_info.rope_scaling['factor'] + if origin_max_model_len is None: + origin_max_model_len = self.model_info.max_model_len + assert origin_max_model_len is not None, '`origin_max_model_len` from model config is not set' + rope_scaling['original_max_position_embeddings'] = origin_max_model_len + + if 'factor' not in rope_scaling: + assert self.max_model_len is not None, ( + 'max_model_len must be set if rope_scaling does not contain a "factor"') + rope_scaling['factor'] = max(float(math.ceil(self.max_model_len / origin_max_model_len)), 1.0) + rope_model_len = int(origin_max_model_len * rope_scaling['factor']) + if self.max_model_len is None: + self.max_model_len = rope_model_len + elif self.max_model_len > rope_model_len: + logger.warning(f'rope config ({rope_model_len} = {rope_scaling["factor"]} * ' + f'{origin_max_model_len}) should be bigger than max_model_len ' + f'from command line ({self.max_model_len})') + self.rope_scaling = rope_scaling + logger.info(f'Setting args.rope_scaling: {rope_scaling}') + logger.info(f'Setting args.max_model_len: {self.max_model_len}') + + def _init_model_info(self) -> torch.dtype: + model_kwargs = self.get_model_kwargs() + if self.tuner_backend == 'unsloth': + model_kwargs['download_model'] = True + self.model_info, self.model_meta = get_model_info_meta(**model_kwargs) + self.task_type = self.model_info.task_type + self.num_labels = self.model_info.num_labels + + self.model_dir = self.model_info.model_dir + self.model_type = self.model_info.model_type + if self.rope_scaling or self.model_info.rope_scaling and self.max_model_len is not None: + self._init_rope_scaling() + return self.model_info.torch_dtype + + def _init_new_special_tokens(self): + if isinstance(self.new_special_tokens, str): + self.new_special_tokens = [self.new_special_tokens] + new_special_tokens = [] + for token in self.new_special_tokens: + if token.endswith('.txt'): + assert os.path.isfile(token), f'special_tokens_path: {token}' + with open(token, 'r', encoding='utf-8') as f: + text = f.read() + new_special_tokens += text.split() + else: + new_special_tokens.append(token) + self.new_special_tokens = new_special_tokens + + def __post_init__(self): + if self.model is None: + raise ValueError(f'Please set --model `, model: {self.model}') + self._init_new_special_tokens() + self.model_suffix = get_model_name(self.model) + self._init_device_map() + self._init_max_memory() + self._init_torch_dtype() + + def get_model_kwargs(self): + return { + 'model_id_or_path': self.model, + 'torch_dtype': self.torch_dtype, + 'model_type': self.model_type, + 'revision': self.model_revision, + 'use_hf': self.use_hf, + 'hub_token': self.hub_token, + 'local_repo_path': self.local_repo_path, + 'device_map': self.device_map, + 'max_memory': self.max_memory, + 'quantization_config': self.get_quantization_config(), + 'attn_impl': self.attn_impl, + 'experts_impl': self.experts_impl, + 'new_special_tokens': self.new_special_tokens, + 'rope_scaling': self.rope_scaling, + 'max_model_len': self.max_model_len, + 'task_type': self.task_type, + 'num_labels': self.num_labels, + 'problem_type': self.problem_type, + 'init_strategy': self.init_strategy, + } diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/quant_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/quant_args.py new file mode 100644 index 0000000000000000000000000000000000000000..790ce7ce5fb6887a8417a403616597ba0458f8c7 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/quant_args.py @@ -0,0 +1,122 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from dataclasses import dataclass +from typing import Literal, Optional + +from swift.model import get_model_processor +from swift.utils import HfConfigFactory, get_modules_to_not_convert + + +@dataclass +class QuantizeArguments: + """A dataclass that holds the configuration for model quantization. + + Args: + quant_method (Optional[str]): The quantization method to use when loading the model. Can be one of {'bnb', + 'hqq', 'eetq', 'quanto', 'fp8'}. Note: This is not required for QLoRA training on pre-quantized AWQ/GPTQ + models. Defaults to None. + quant_bits (Optional[Union[int, str]]): The number of bits for quantization, e.g., {1, 2, 3, 4, 8, 'float8'}. + Defaults to None. + hqq_axis (Optional[int]): The quantization axis for HQQ quantization. Defaults to None. + bnb_4bit_compute_dtype (Optional[str]): The compute data type for 4-bit BNB quantization. Can be one of { + 'float16', 'bfloat16', 'float32'}. Defaults to None, which will use the model's `torch_dtype`. + bnb_4bit_quant_type (str): The quantization type for 4-bit BNB quantization. Can be one of {'fp4', 'nf4'}. + Defaults to 'nf4'. + bnb_4bit_use_double_quant (bool): Whether to use double quantization for 4-bit BNB quantization. + Defaults to True. + bnb_4bit_quant_storage (Optional[str]): The storage type for packing quantized 4-bit parameters in BNB. + Defaults to None. + """ + # awq, gptq, and aqlm need to be pre-quantized models. + # It can be detected automatically, without the need to pass in. + # while bnb, hqq, and eetq can be quantized during SFT using the original models. + quant_method: Literal['bnb', 'hqq', 'eetq', 'quanto', 'fp8'] = None + # bnb: 4,8; hqq: 1,2,3,4,8'; eetq: 8 + # awq: 4; gptq: 2,3,4,8 + quant_bits: Literal[1, 2, 3, 4, 8, 'float8'] = None + # hqq + hqq_axis: Optional[int] = None + # bnb + bnb_4bit_compute_dtype: Literal['float16', 'bfloat16', 'float32', None] = None + bnb_4bit_quant_type: Literal['fp4', 'nf4'] = 'nf4' + bnb_4bit_use_double_quant: bool = True + bnb_4bit_quant_storage: Optional[str] = None + + def get_quantization_config(self): + if self.quant_method is None or self.quant_method in {'awq', 'gptq', 'gptq_v2'}: + return None + assert self.quant_method in {'bnb', 'hqq', 'eetq', 'quanto', 'fp8'} + if self.quant_method != 'fp8' and self.quant_bits is None: + raise ValueError(f'Please set the quant_bits. args.quant_bits: {self.quant_bits}') + if self.quant_method == 'bnb': + if self.quant_bits == 4: + load_in_4bit, load_in_8bit = True, False + elif self.quant_bits == 8: + load_in_4bit, load_in_8bit = False, True + else: + raise ValueError(f'bnb not support quant_bits: {self.quant_bits}') + + from transformers import BitsAndBytesConfig + llm_int8_skip_modules = self.get_modules_to_not_convert() + quantization_config = BitsAndBytesConfig( + load_in_4bit=load_in_4bit, + load_in_8bit=load_in_8bit, + bnb_4bit_compute_dtype=self.bnb_4bit_compute_dtype, + bnb_4bit_quant_type=self.bnb_4bit_quant_type, + bnb_4bit_use_double_quant=self.bnb_4bit_use_double_quant, + bnb_4bit_quant_storage=self.bnb_4bit_quant_storage, + llm_int8_skip_modules=llm_int8_skip_modules) + elif self.quant_method == 'fp8': + if not hasattr(self, 'model_info'): + return + from transformers import FineGrainedFP8Config + with torch.device('meta'): + hf_model, _ = get_model_processor(self.model_dir, model_type=self.model_type, return_dummy_model=True) + modules_to_not_convert = get_modules_to_not_convert(hf_model) + quantization_config = FineGrainedFP8Config(modules_to_not_convert=modules_to_not_convert) + elif self.quant_method == 'hqq': + from transformers import HqqConfig + quantization_config = HqqConfig(nbits=self.quant_bits, axis=self.hqq_axis) + elif self.quant_method == 'quanto': + from transformers import QuantoConfig + if self.quant_bits == 8: + weights = 'int8' + elif self.quant_bits == 'float8': + weights = 'float8' + elif self.quant_bits == 4: + weights = 'int4' + elif self.quant_bits == 2: + weights = 'int2' + else: + raise ValueError('quanto quantization only support quant bits 2/4/8/float8') + quantization_config = QuantoConfig(weights=weights) + else: # 'eetq' + from transformers import EetqConfig + quantization_config = EetqConfig(f'int{self.quant_bits}') + + return quantization_config + + def get_modules_to_not_convert(self): + if not hasattr(self, 'model_meta') or not hasattr(self, 'model_info'): + return None + model_arch = self.model_meta.model_arch + res = [] + if self.model_info.is_moe_model: + res += ['mlp.gate', 'mlp.shared_expert_gate'] + if model_arch is not None: + for key in ['vision_tower', 'aligner']: + value = getattr(model_arch, key, None) + if value: + res += value + if not res: + return None + res.append('lm_head') + return res + + def __post_init__(self): + if self.bnb_4bit_compute_dtype is None: + if self.torch_dtype in {torch.float16, torch.float32}: + self.bnb_4bit_compute_dtype = torch.float32 + elif self.torch_dtype == torch.bfloat16: + self.bnb_4bit_compute_dtype = torch.bfloat16 + self.bnb_4bit_compute_dtype: torch.dtype = HfConfigFactory.to_torch_dtype(self.bnb_4bit_compute_dtype) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/template_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/template_args.py new file mode 100644 index 0000000000000000000000000000000000000000..a2a711aca59961a8837309fc01d32c5d431a7ec1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/base_args/template_args.py @@ -0,0 +1,207 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +from dataclasses import dataclass, field +from typing import Literal, Optional + +from swift.template import TEMPLATE_MAPPING, get_template_meta +from swift.utils import get_logger + +logger = get_logger() + + +@dataclass +class TemplateArguments: + """TemplateArguments class holds various arguments for template configuration. + + This dataclass manages settings related to how data is formatted and processed using templates, including + tokenization, truncation, loss calculation, and special handling for multimodal and agent-based models. + + Args: + template (Optional[str]): The dialogue template type. Defaults to None, which automatically selects the + template corresponding to the model type. Refer to the list of supported models for mappings. + system (Optional[str]): Custom system prompt. Can be a string or a path to a .txt file. Defaults to None, + which uses the default system from the registered template. + Note: The priority for the system prompt is as follows: + 1. System prompt from the dataset. + 2. The `--system` command-line argument. + 3. The `default_system` set when the template was registered. + max_length (Optional[int]): The maximum number of tokens for a single sample after tokenization. Samples + exceeding this length are handled according to `truncation_strategy` to prevent OOM errors. Defaults to + None, which uses the model's maximum supported length (`max_model_len`). In PPO, GRPO, and inference + scenarios, this argument specifies the `max_prompt_length`. + truncation_strategy (Literal['delete', 'left', 'right', 'split']): Strategy for handling samples exceeding + `max_length`. Options are 'delete', 'left' (truncate from the left), 'right' (truncate from the right), + and 'split' (split into multiple samples). Defaults to 'delete'. + Note: The 'split' strategy is only supported during pre-training (e.g., `swift/megatron pt`), + and is incompatible with `cached_dataset`. It splits long samples to avoid wasting tokens. + Note: For multimodal models, setting this to 'left' or 'right' preserves all image tokens, which may lead + to OOM errors. + max_pixels (Optional[int]): The maximum number of pixels (H*W) for an input image in a multimodal model. + Images exceeding this limit will be scaled down to prevent OOM errors. Defaults to None, meaning no limit. + Note: This parameter applies to all multimodal models. The model-specific `MAX_PIXELS` parameter for + Qwen2.5-VL is separate and only applies to that model. + agent_template (Optional[str]): The Agent template to use. This determines how the 'tools' list is converted + into a 'system' prompt, how tool calls are extracted from the model's response during inference, and the + format for tool call messages. Options include "react_en", "hermes", "glm4", "qwen_en", "toolbench", etc. + Defaults to None, which auto-selects based on the model type. Refer to the Agent documentation for more + details. + norm_bbox (Optional[Literal['norm1000', 'none']]): Controls how bounding box coordinates (from the "bbox" + field in the dataset) are scaled. 'norm1000' scales coordinates to a 1000x1000 grid, while 'none' performs + no scaling. Defaults to None, which auto-selects based on the model. This handles cases where images are + resized during training (e.g., due to `max_pixels`). + use_chat_template (bool): Whether to use the chat template or the generation template. The generation template + is typically used for pre-training. Defaults to True. + Note: Defaults to False for `swift pt`, which uses the generation template. This parameter is compatible + with multimodal models. + padding_side (Literal['left', 'right']): The side to pad on when `batch_size >= 2` during training. + Options are 'left' or 'right'. Defaults to 'right'. For inference with `batch_size >= 2`, padding is always + on the left. + Note: Defaults to 'left' for PPO and GKD. + padding_free (bool): If True, flattens the data within a batch to avoid padding, reducing memory usage and + speeding up training. Sequences within the batch remain causally isolated. Defaults to False. Supported for + CPT/SFT/DPO/GRPO/KTO/GKD. + Note: This requires `--attn_impl flash_attn` and `transformers>=4.44`. Compared to packing, padding_free + has no preprocessing overhead, but packing offers faster training speeds and more stable memory usage. + loss_scale (str): Loss weight configuration for training tokens. Default is `'default'`. + loss_scale includes 3 basic strategies: 'default', 'last_round', 'all', and other strategies: + 'ignore_empty_think', 'ignore_think_prefix', and agent-specific ones: 'react', 'hermes', 'qwen', + 'agentflan', 'alpha_umi', etc. + For available options, refer to + [loss_scale module](https://github.com/modelscope/ms-swift/blob/main/swift/loss_scale/mapping.py). + ms-swift supports mixing basic strategies with other strategies, + for example: `'default+ignore_empty_think'`, `'last_round+ignore_empty_think'`. + If no basic strategy is specified, it defaults to 'default', + for example: 'hermes' is equivalent to 'default+hermes'. + Multiple non-base strategies can be chained together + (each strategy processes the output segments of the previous one, with weights + multiplied accordingly). For example: `'last_round+hermes+ignore_empty_think'`, where + `'last_round'` is the base strategy, and `'hermes+ignore_empty_think'` represents a + chain of multiple non-base strategies that share the same base strategy. + - 'default': All responses (including history) are calculated with weight 1 for cross-entropy loss + (**system/user/multimodal tokens in messages and `tool_response` parts in Agent training are + not included in loss calculation**). (**Default value for SFT**) + - 'last_round': Only calculate loss for the last round response. The last round + means all content after the last "user". (**Default value for RLHF**) + - 'all': Calculate loss for all tokens. (**Default value for `swift pt`**) + - 'ignore_empty_think': Ignore loss computation for empty `'\n\n\n\n'` + (as long as it matches the regex `'\\s*\\s*'`). + - 'ignore_think_prefix': Ignore loss computation for `'\n'` at the start of a response + (or `''` when no newline follows). The reasoning content and `''` remain trainable. + - 'react', 'hermes', 'qwen': Adjust the loss weight of the `tool_call` part to 2. + sequence_parallel_size (int): The size of sequence parallelism. Defaults to 1. Currently supported for CPT, + SFT, DPO, and GRPO. + template_backend (Literal['swift', 'jinja']): The backend to use for templating. Options are 'swift' or + 'jinja'. Defaults to 'swift'. If 'jinja' is used, it will leverage `transformers.apply_chat_template`. + Note: The 'jinja' backend is only supported for inference, not for training, as it cannot determine the + token range for loss calculation. + response_prefix (Optional[str]): A prefix string for the response, e.g., '\\n' for Qwen-32B. This + parameter only affects inference. Defaults to None, which is auto-set based on the model. + enable_thinking (Optional[bool]): This parameter takes effect during inference, + indicating whether to enable thinking mode. Default is None, the default value is determined by the + template (model) type (True for thinking/hybrid thinking templates, False for non-thinking templates). + If enable_thinking is False, a non-thinking prefix is added, for example the Qwen3-8B hybrid thinking + model adds the prefix `'\n\n\n\n'`, while Qwen3-8B-Thinking does not add a prefix. + If enable_thinking is True, a thinking prefix is added, for example `'\n'`. + Note: The priority of this parameter is lower than the response_prefix parameter. + preserve_thinking (Optional[bool]): Whether to preserve historical thinking content during inference and + training. When set to `True`, thinking content from all rounds is retained. When set to `False`, + only the thinking content from the last round is retained (i.e., the content following the last + user message). Defaults to `None`. + Default behavior: For thinking models (thinking/hybrid-thinking) or when `enable_thinking` is + explicitly enabled, this is set to `False` by default during inference and training, retaining + only the last round of thinking content. If the `loss_scale` base strategy during training is + not `'last_round'` (e.g., `'default'`), it defaults to `True`, and historical thinking content will + not be removed. + add_non_thinking_prefix (bool): This parameter only takes effect during training, indicating whether to + add a non-thinking prefix to data samples whose assistant part does not start with the thinking + marker `''` (typically hybrid thinking models contain a non-thinking prefix). + This feature allows swift's built-in datasets to train hybrid thinking models. Default value is True. + For example: the non-thinking prefix for the Qwen3-8B hybrid thinking model is + `'\n\n\n\n'`, while the non-thinking prefix for Qwen3-8B-Thinking/Instruct is `''`. + Note: During training, if the basic strategy of loss_scale is last_round, this modification is only + applied to the last round; otherwise, for example 'default' or 'all', this modification is applied to + every round of data. If set to False, no non-thinking prefix is added to data samples. + + + """ + template: Optional[str] = field( + default=None, metadata={'help': f'template choices: {list(TEMPLATE_MAPPING.keys())}'}) + system: Optional[str] = None # Override the default_system in the template. + max_length: Optional[int] = None + + truncation_strategy: Literal['delete', 'left', 'right', 'split', None] = None + max_pixels: Optional[int] = None + agent_template: Optional[str] = None + norm_bbox: Literal['norm1000', 'none', None] = None + use_chat_template: Optional[bool] = None + padding_side: Literal['left', 'right'] = 'right' + # train + padding_free: bool = False + loss_scale: str = 'default' + sequence_parallel_size: int = 1 + is_binary_loss_scale: Optional[bool] = None + # infer/deploy + template_backend: Literal['swift', 'jinja'] = 'swift' + # thinking + response_prefix: Optional[str] = None + enable_thinking: Optional[bool] = None + preserve_thinking: Optional[bool] = None + add_non_thinking_prefix: bool = True + disable_ignore_empty_think: bool = False + + def __post_init__(self): + if getattr(self, 'model_meta', None) is not None: + self.template_meta = get_template_meta(self.model_info, self.model_meta, template_type=self.template) + self.template = self.template_meta.template_type + if self.use_chat_template is None: + self.use_chat_template = True + if self.system is not None: + if self.system.endswith('.txt'): + assert os.path.isfile(self.system), f'self.system: {self.system}' + with open(self.system, 'r', encoding='utf-8') as f: + self.system = f.read() + else: + self.system = self.system.replace('\\n', '\n') + if self.response_prefix is not None: + self.response_prefix = self.response_prefix.replace('\\n', '\n') + if self.truncation_strategy is None: + self.truncation_strategy = 'delete' + self._set_loss_scale() + + def _set_loss_scale(self): + """For hybrid thinking models, automatically append '+ignore_empty_think' to loss_scale.""" + if not self.disable_ignore_empty_think and getattr(self, 'template_meta', None) is not None: + template_meta = self.template_meta + if template_meta.is_thinking and template_meta.non_thinking_prefix: + # hybrid thinking model detected + if self.loss_scale and 'ignore_empty_think' not in self.loss_scale: + self.loss_scale = self.loss_scale + '+ignore_empty_think' + + def get_template_kwargs(self): + truncation_strategy = self.truncation_strategy + if truncation_strategy == 'delete': + truncation_strategy = 'raise' + return { + 'template_type': self.template, + 'default_system': self.system, + 'max_length': self.max_length, + 'truncation_strategy': truncation_strategy, + 'max_pixels': self.max_pixels, + 'agent_template': self.agent_template, + 'norm_bbox': self.norm_bbox, + 'use_chat_template': self.use_chat_template, + 'remove_unused_columns': self.remove_unused_columns, # from DataArguments + 'padding_side': self.padding_side, + # train + 'padding_free': self.padding_free, + 'loss_scale': self.loss_scale, + 'is_binary_loss_scale': self.is_binary_loss_scale, + 'sequence_parallel_size': self.sequence_parallel_size, + # infer/deploy + 'template_backend': self.template_backend, + # thinking + 'response_prefix': self.response_prefix, + 'enable_thinking': self.enable_thinking, + 'preserve_thinking': self.preserve_thinking, + 'add_non_thinking_prefix': self.add_non_thinking_prefix, + } diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/deploy_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/deploy_args.py new file mode 100644 index 0000000000000000000000000000000000000000..5960b6516249ca602ad0b916140d439a29505e19 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/deploy_args.py @@ -0,0 +1,180 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +from dataclasses import dataclass +from typing import Literal, Optional + +from swift.utils import find_free_port, get_device_count, get_logger, safe_snapshot_download +from .base_args import BaseArguments +from .infer_args import InferArguments + +logger = get_logger() + + +@dataclass +class DeployArguments(InferArguments): + """Arguments for model deployment. + + This dataclass, which extends InferArguments, is used to define the arguments required for deploying a model. + + Args: + host (str): The host address to bind the server to. Defaults to '0.0.0.0'. + port (int): The port number to bind the server to. Defaults to 8000. + api_key (Optional[str]): The API key for authentication. Defaults to None. + ssl_keyfile (Optional[str]): The path to the SSL key file. Defaults to None. + ssl_certfile (Optional[str]): The path to the SSL certificate file. Defaults to None. + owned_by (str): The owner of the deployment. Defaults to 'swift'. + served_model_name (Optional[str]): The name of the model being served. If None, the model's suffix is used by + default. + verbose (bool): Whether to log detailed request information. Defaults to True. + Note: This defaults to False when used in 'swift app' or 'swift eval'. + log_interval (int): The interval in seconds for printing tokens/s statistics. Set to -1 to disable. Defaults + to 20. + log_level (Literal['critical', 'error', 'warning', 'info', 'debug', 'trace']): Log level. Defaults to 'info'. + max_logprobs (int): The maximum number of logprobs to return to the client. Defaults to 20. + vllm_use_async_engine (Optional[bool]): Whether to use async engine for vLLM.If not set, it defaults to `True` + for deployment scenarios. + """ + host: str = '0.0.0.0' + port: int = 8000 + api_key: Optional[str] = None + ssl_keyfile: Optional[str] = None + ssl_certfile: Optional[str] = None + + owned_by: str = 'swift' + served_model_name: Optional[str] = None + verbose: bool = True # Whether to log request_info + log_interval: int = 20 # Interval for printing global statistics + log_level: Literal['critical', 'error', 'warning', 'info', 'debug', 'trace'] = 'info' + + max_logprobs: int = 20 + vllm_use_async_engine: Optional[bool] = None + + def __post_init__(self): + # default to True for deployment scenarios + if self.vllm_use_async_engine is None: + self.vllm_use_async_engine = True + super().__post_init__() + self.port = find_free_port(self.port) + + def _init_adapters(self): + if isinstance(self.adapters, str): + self.adapters = [self.adapters] + self.adapter_mapping = {} + adapters = [] + for i, adapter in enumerate(self.adapters): + adapter_path = adapter.split('=') + if len(adapter_path) == 1: + adapter_path = (None, adapter_path[0]) + adapter_name, adapter_path = adapter_path + adapter_path = safe_snapshot_download(adapter_path, use_hf=self.use_hf, hub_token=self.hub_token) + if adapter_name is None: + adapters.append(adapter_path) + else: + self.adapter_mapping[adapter_name] = adapter_path + self.adapters = adapters + + def _init_ckpt_dir(self, adapters=None): + return super()._init_ckpt_dir(self.adapters + list(self.adapter_mapping.values())) + + def _init_stream(self): + return BaseArguments._init_stream(self) + + +@dataclass +class RolloutArguments(DeployArguments): + """Arguments for the Rollout phase in online/reinforcement learning. + + This dataclass inherits from DeployArguments and adds specific parameters for the Rollout process in online + learning, such as GRPO. + + Args: + multi_turn_scheduler (Optional[str]): The scheduler for multi-turn GRPO training. Pass the name of the + corresponding plugin implemented in `swift/rollout/multi_turn.py`. Defaults to None. Refer to the + documentation for details. + max_turns (Optional[int]): The maximum number of turns in multi-turn GRPO training. If None, no limit is + imposed. Defaults to None. + vllm_enable_lora (bool): Whether to enable the vLLM Engine to load LoRA adapters. Enabling this can accelerate + weight synchronization during LoRA training. Defaults to False. Refer to the documentation for details. + vllm_max_lora_rank (int): The LoRA rank parameter for the vLLM Engine. This value must be greater than or + equal to the `lora_rank` used for training; setting them as equal is recommended. Defaults to 16. + """ + vllm_use_async_engine: Optional[bool] = None + use_gym_env: Optional[bool] = None + # only for GRPO rollout with AsyncEngine, see details in swift/rollout/multi_turn + multi_turn_scheduler: Optional[str] = None + max_turns: Optional[int] = None + vllm_enable_lora: bool = False + vllm_max_lora_rank: int = 16 + # GYM env + gym_env: Optional[str] = None + context_manager: Optional[str] = None + + def __post_init__(self): + self._set_default_engine_type() + super().__post_init__() + self._check_args() + self._check_device_count() + self._check_vllm_enable_expert_parallel() + self._check_deprecated_args() + self._set_default_audio_load_backend() + + def _set_default_engine_type(self): + if self.vllm_use_async_engine is None: + if self.multi_turn_scheduler: + self.vllm_use_async_engine = True + else: + self.vllm_use_async_engine = False + + if self.use_gym_env is None: + self.use_gym_env = self.gym_env is not None + + def _check_args(self): + if self.vllm_pipeline_parallel_size > 1: + raise ValueError('RolloutArguments does not support pipeline parallelism, ' + 'please set vllm_pipeline_parallel_size to 1.') + + if self.vllm_reasoning_parser is not None: + raise ValueError('vllm_reasoning_parser is not supported for Rollout, please unset it.') + + if self.multi_turn_scheduler and not self.vllm_use_async_engine: + raise ValueError('please set vllm_use_async_engine to True with multi-turn scheduler.') + + def _check_device_count(self): + local_device_count = get_device_count() + required_device_count = self.vllm_data_parallel_size * self.vllm_tensor_parallel_size + + if local_device_count < required_device_count: + msg = (f'Error: local_device_count ({local_device_count}) must be greater than or equal to ' + f'the product of vllm_data_parallel_size ({self.vllm_data_parallel_size}) and ' + f'vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}). ' + f'Current required_device_count = {required_device_count}.') + raise ValueError(msg) + + if local_device_count > required_device_count: + logger.warning_once( + f'local_device_count ({local_device_count}) is greater than required_device_count ({required_device_count}). ' # noqa + f'Only the first {required_device_count} devices will be utilized for rollout. ' + f'To fully utilize resources, set vllm_tensor_parallel_size * vllm_data_parallel_size = device_count. ' # noqa + f'device_count: {local_device_count}, ' + f'vllm_tensor_parallel_size: {self.vllm_tensor_parallel_size}, ' + f'vllm_data_parallel_size: {self.vllm_data_parallel_size}, ' + f'required_device_count: {required_device_count}.') + + def _check_vllm_enable_expert_parallel(self): + if self.vllm_enable_expert_parallel and not self.vllm_use_async_engine: + self.vllm_use_async_engine = True + logger.warning('vllm_enable_expert_parallel is only supported with vllm_use_async_engine, ' + 'set vllm_use_async_engine to True.') + + def _check_deprecated_args(self): + if self.context_manager is not None: + raise ValueError('The "context_manager" argument has been removed. ' + 'If you need to dynamically modify the conversation history between rollout turns ' + '(e.g. history compression, prompt injection), implement that logic in a custom ' + '`MultiTurnScheduler` subclass by overriding `step` / `run`, ' + 'and pass it via `--multi_turn_scheduler your_scheduler_name`.') + + def _set_default_audio_load_backend(self): + # Rollout uses GRPOVllmEngine (vLLM-only); align audio decode with vLLM multimodal loader. + if os.getenv('SWIFT_AUDIO_LOAD_BACKEND') is None: + os.environ['SWIFT_AUDIO_LOAD_BACKEND'] = 'soundfile_pyav' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/eval_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/eval_args.py new file mode 100644 index 0000000000000000000000000000000000000000..64c692b6eec07ec3c052b57e05846c2244693404 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/eval_args.py @@ -0,0 +1,130 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import datetime as dt +import os +from dataclasses import dataclass, field +from typing import List, Literal, Optional, Union + +from swift.model import get_matched_model_meta +from swift.utils import get_logger, json_parse_to_dict, to_abspath +from .deploy_args import DeployArguments + +logger = get_logger() + + +@dataclass +class EvalArguments(DeployArguments): + """A dataclass that extends DeployArguments to define model evaluation arguments. + + These arguments control the evaluation process, including the choice of backend, datasets, generation parameters, + and other configurations. + + Args: + eval_dataset (List[str]): List of evaluation datasets. Please refer to the evaluation documentation for + available options. Defaults to []. + eval_limit (Optional[int]): The number of samples to take from each evaluation dataset. If None, all samples + are used. Defaults to None. + eval_dataset_args (Optional[Union[Dict, str]]): Evaluation dataset parameters, in JSON format, can be set for + multiple datasets. Defaults to None. + eval_generation_config (Optional[Union[Dict, str]]): The model's inference configuration for evaluation, + provided as a JSON string (e.g., '{"max_new_tokens": 512}'). Defaults to None. + eval_output_dir (str): The directory to store evaluation results. Defaults to 'eval_output'. + eval_backend (str): The evaluation backend. Can be 'Native', 'OpenCompass', or 'VLMEvalKit'. Defaults to + 'Native'. + local_dataset (bool): Whether to automatically download extra datasets required for certain evaluations + (e.g., CMB). If True, a 'data' folder will be created in the current directory for the datasets. This + download occurs only once, and subsequent runs will use the cache. Defaults to False. + Note: By default, evaluation uses datasets from `~/.cache/opencompass`. When this is set to True, the + `data` folder in the current directory is used instead. + temperature (float): The temperature for sampling, which overrides the default generation config. Defaults + to 0.0. + verbose (bool): Whether to output verbose information during the evaluation process. Defaults to False. + eval_num_proc (int): The maximum number of concurrent clients for evaluation. Defaults to 16. + extra_eval_args (Optional[Union[Dict, str]]): Additional evaluation arguments, provided as a JSON string. + These are only effective when using the 'Native' backend. Refer to the documentation for more details on + available arguments. Defaults to {}. + eval_url (Optional[str]): The URL for the evaluation service (e.g., 'http://localhost:8000/v1'). If not + specified, evaluation runs on the locally deployed model. See documentation for more examples. Defaults + to None. + """ + eval_dataset: List[str] = field(default_factory=list) + eval_limit: Optional[int] = None + eval_dataset_args: Optional[Union[dict, str]] = None + eval_generation_config: Optional[Union[dict, str]] = None + eval_output_dir: str = 'eval_output' + eval_backend: Literal['Native', 'OpenCompass', 'VLMEvalKit'] = 'Native' + local_dataset: bool = False + + temperature: Optional[float] = 0. + verbose: bool = False + eval_num_proc: int = 16 + extra_eval_args: Optional[Union[dict, str]] = field(default_factory=dict) + # If eval_url is set, ms-swift will not perform deployment operations and + # will directly use the URL for evaluation. + eval_url: Optional[str] = None + + def __post_init__(self): + super().__post_init__() + self._init_eval_url() + self._init_eval_dataset() + self.eval_dataset_args = json_parse_to_dict(self.eval_dataset_args) + self.eval_generation_config = json_parse_to_dict(self.eval_generation_config) + self.extra_eval_args = json_parse_to_dict(self.extra_eval_args) + self.eval_output_dir = to_abspath(self.eval_output_dir) + logger.info(f'eval_output_dir: {self.eval_output_dir}') + + def _init_eval_url(self): + # [compat] + if self.eval_url and 'chat/completions' in self.eval_url: + self.eval_url = self.eval_url.split('/chat/completions', 1)[0] + + @staticmethod + def list_eval_dataset(eval_backend=None): + from evalscope.api.registry import BENCHMARK_REGISTRY + from evalscope.backend.opencompass import OpenCompassBackendManager + from evalscope.constants import EvalBackend + res = { + EvalBackend.NATIVE: list(sorted(BENCHMARK_REGISTRY.keys())), + EvalBackend.OPEN_COMPASS: sorted(OpenCompassBackendManager.list_datasets()), + } + try: + from evalscope.backend.vlm_eval_kit import VLMEvalKitBackendManager + vlm_datasets = VLMEvalKitBackendManager.list_supported_datasets() + res[EvalBackend.VLM_EVAL_KIT] = sorted(vlm_datasets) + except ImportError: + # fix cv2 import error + if eval_backend == 'VLMEvalKit': + raise + return res + + def _init_eval_dataset(self): + if isinstance(self.eval_dataset, str): + self.eval_dataset = [self.eval_dataset] + + all_eval_dataset = self.list_eval_dataset(self.eval_backend) + dataset_mapping = {dataset.lower(): dataset for dataset in all_eval_dataset[self.eval_backend]} + valid_dataset = [] + for dataset in self.eval_dataset: + if dataset.lower() not in dataset_mapping: + raise ValueError( + f'eval_dataset: {dataset} is not supported.\n' + f'eval_backend: {self.eval_backend} supported datasets: {all_eval_dataset[self.eval_backend]}') + valid_dataset.append(dataset_mapping[dataset.lower()]) + self.eval_dataset = valid_dataset + + logger.info(f'eval_backend: {self.eval_backend}') + logger.info(f'eval_dataset: {self.eval_dataset}') + + def _init_result_path(self, folder_name: str) -> None: + self.time = dt.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f') + result_dir = self.ckpt_dir or f'result/{self.model_suffix}' + os.makedirs(result_dir, exist_ok=True) + self.result_jsonl = to_abspath(os.path.join(result_dir, 'eval_result.jsonl')) + if not self.eval_url: + super()._init_result_path('eval_result') + + def _init_torch_dtype(self) -> None: + if self.eval_url: + self.model_meta = get_matched_model_meta(self.model) + self.model_info = None + return + super()._init_torch_dtype() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/export_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/export_args.py new file mode 100644 index 0000000000000000000000000000000000000000..ae82949afe800ae582905bf4f5fdea892e290cae --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/export_args.py @@ -0,0 +1,154 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import torch +import torch.distributed as dist +from dataclasses import dataclass +from typing import Literal, Optional + +from swift.utils import HfConfigFactory, get_logger, init_process_group, set_default_ddp_config, to_abspath +from .base_args import BaseArguments +from .merge_args import MergeArguments + +logger = get_logger() + + +@dataclass +class ExportArguments(MergeArguments, BaseArguments): + """ExportArguments is a dataclass that inherits from BaseArguments and MergeArguments. + + Args: + output_dir (Optional[str]): Directory to save the exported results. Defaults to None, which automatically sets + a path with an appropriate suffix. + quant_method (Optional[str]): The quantization method. Can be 'awq', 'gptq', 'bnb', 'fp8', or 'gptq_v2'. + Defaults to None. See examples for more details. + quant_n_samples (int): Number of samples for GPTQ/AWQ calibration. Defaults to 256. + quant_batch_size (int): The batch size for quantization. Defaults to 1. + group_size (int): The group size for quantization. Defaults to 128. + to_cached_dataset (bool): Whether to tokenize and export the dataset in advance as a cached dataset. Defaults + to False. Note: You can specify the validation set content through + `--split_dataset_ratio` or `--val_dataset`. + to_ollama (bool): Whether to generate the `Modelfile` required by Ollama. Defaults to False. + to_mcore (bool): Whether to convert Hugging Face format weights to Megatron-Core format. Defaults to False. + to_hf (bool): Whether to convert Megatron-Core format weights to Hugging Face format. Defaults to False. + mcore_model (Optional[str]): The path to the Megatron-Core format model. Defaults to None. + mcore_adapter (Optional[str]): A list of adapter paths for the Megatron-Core format model. Defaults to []. + thread_count (Optional[int]): The number of model shards when `to_mcore` is True. Defaults to None, which + automatically sets the number based on the model size to keep the largest shard under 10GB. + test_convert_precision (bool): Whether to test the precision error of weight conversion between Hugging Face + and Megatron-Core formats. Defaults to False. + test_convert_dtype (str): The dtype to use for the conversion precision test. Defaults to 'float32'. + save_missing_weights (bool): Whether to copy weights that exist in the source checkpoint but are not + supported by Megatron-Core (e.g. the `mtp.*` DSpark weights of DeepSeek-V4-Flash) into the exported + model, so that they are not lost. Defaults to False. + push_to_hub (bool): Whether to push the output to the Model Hub. Defaults to False. See examples for more + details. + hub_model_id (Optional[str]): The model ID for pushing to the Hub (e.g., 'user_name/repo_name' or 'repo_name'). + Defaults to None. + hub_private_repo (bool): Whether the Hub repository is private. Defaults to False. + commit_message (str): The commit message for pushing to the Hub. Defaults to 'update files'. + to_peft_format (bool): Whether to export in PEFT format. This argument is for compatibility and currently has + no effect. Defaults to False. + exist_ok (bool): If the output_dir exists, do not raise an exception and overwrite its contents. Defaults to + False. + """ + output_dir: Optional[str] = None + + # awq/gptq + quant_method: Literal['awq', 'gptq', 'bnb', 'fp8', 'gptq_v2'] = None + quant_n_samples: int = 256 + quant_batch_size: int = 1 + group_size: int = 128 + + # cached_dataset + to_cached_dataset: bool = False + template_mode: Literal['train', 'rlhf', 'kto'] = 'train' + + # ollama + to_ollama: bool = False + + # megatron + to_mcore: bool = False + to_hf: bool = False + mcore_model: Optional[str] = None + mcore_adapter: Optional[str] = None + thread_count: Optional[int] = None + test_convert_precision: bool = False + test_convert_dtype: str = 'float32' + save_missing_weights: bool = False + + # push to ms hub + push_to_hub: bool = False + # 'user_name/repo_name' or 'repo_name' + hub_model_id: Optional[str] = None + hub_private_repo: bool = False + commit_message: str = 'update files' + # compat + to_peft_format: bool = False + exist_ok: bool = False + + def load_args_from_ckpt(self) -> None: + if self.to_cached_dataset: + return + super().load_args_from_ckpt() + + def _init_output_dir(self): + if self.output_dir is None: + ckpt_dir = self.ckpt_dir or f'./{self.model_suffix}' + ckpt_dir, ckpt_name = os.path.split(ckpt_dir) + if self.to_peft_format: + suffix = 'peft' + elif self.quant_method: + suffix = f'{self.quant_method}' + if self.quant_bits is not None: + suffix += f'-int{self.quant_bits}' + elif self.to_ollama: + suffix = 'ollama' + elif self.merge_lora: + suffix = 'merged' + elif self.to_mcore: + suffix = 'mcore' + elif self.to_hf: + suffix = 'hf' + elif self.to_cached_dataset: + suffix = 'cached_dataset' + else: + return + + self.output_dir = os.path.join(ckpt_dir, f'{ckpt_name}-{suffix}') + + self.output_dir = to_abspath(self.output_dir) + if not self.exist_ok and os.path.exists(self.output_dir): + raise FileExistsError(f'args.output_dir: `{self.output_dir}` already exists.') + logger.info(f'args.output_dir: `{self.output_dir}`') + + def __post_init__(self): + if self.quant_batch_size == -1: + self.quant_batch_size = None + if self.quant_bits and self.quant_method is None: + raise ValueError('Please specify the quantization method using `--quant_method awq/gptq/bnb`.') + if self.quant_method and self.quant_bits is None and self.quant_method != 'fp8': + raise ValueError('Please specify `--quant_bits`.') + if self.quant_method in {'gptq', 'awq'} and self.torch_dtype is None: + self.torch_dtype = torch.float16 + if self.to_mcore or self.to_hf: + if self.merge_lora: + self.merge_lora = False + logger.warning('`swift export --to_mcore/to_hf` does not support the `--merge_lora` parameter. ' + 'To export LoRA delta weights, please use `megatron export`') + + self.mcore_model = to_abspath(self.mcore_model, check_path_exist=True) + if not dist.is_initialized(): + set_default_ddp_config() + init_process_group(backend=self.ddp_backend, timeout=self.ddp_timeout) + + BaseArguments.__post_init__(self) + self._init_output_dir() + self.test_convert_dtype = HfConfigFactory.to_torch_dtype(self.test_convert_dtype) + if self.quant_method in {'gptq', 'awq'} and len(self.dataset) == 0: + raise ValueError(f'self.dataset: {self.dataset}, Please input the quant dataset.') + if self.to_cached_dataset: + self.lazy_tokenize = False + if self.packing: + raise ValueError('Packing will be handled during training; here we only perform tokenization ' + 'in advance, so you do not need to set up packing separately.') + assert not self.streaming, 'not supported' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/infer_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/infer_args.py new file mode 100644 index 0000000000000000000000000000000000000000..08e65ea8a011951eee37c23ea311e15a697cf559 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/infer_args.py @@ -0,0 +1,237 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import datetime as dt +import os +import torch.distributed as dist +from dataclasses import dataclass +from typing import Literal, Optional + +from swift.rlhf_trainers import VllmArguments +from swift.utils import get_logger, init_process_group, is_dist, to_abspath +from .base_args import BaseArguments +from .merge_args import MergeArguments + +logger = get_logger() + + +@dataclass +class LmdeployArguments: + """Holds the configuration arguments for lmdeploy. + + Args: + lmdeploy_tp (int): The tensor parallelism size. Defaults to 1. + lmdeploy_session_len (Optional[int]): The maximum session length. Defaults to None. + lmdeploy_cache_max_entry_count (float): The percentage of GPU memory to be used by the K/V cache. Defaults + to 0.8. + lmdeploy_quant_policy (int): The quantization policy for the K/V cache. Set to 4 or 8 for 4-bit or 8-bit + quantization respectively. Defaults to 0, which means no quantization. + lmdeploy_vision_batch_size (int): The `max_batch_size` parameter to be passed to `VisionConfig`. Defaults to 1. + """ + + # lmdeploy + lmdeploy_tp: int = 1 + lmdeploy_session_len: Optional[int] = None + lmdeploy_cache_max_entry_count: float = 0.8 + lmdeploy_quant_policy: int = 0 # e.g. 4, 8 + lmdeploy_vision_batch_size: int = 1 # max_batch_size in VisionConfig + + def get_lmdeploy_engine_kwargs(self): + kwargs = { + 'tp': self.lmdeploy_tp, + 'session_len': self.lmdeploy_session_len, + 'cache_max_entry_count': self.lmdeploy_cache_max_entry_count, + 'quant_policy': self.lmdeploy_quant_policy, + 'vision_batch_size': self.lmdeploy_vision_batch_size + } + if dist.is_initialized(): + kwargs.update({'devices': [dist.get_rank()]}) + return kwargs + + +@dataclass +class SglangArguments: + """Arguments for configuring the SGLang backend. + + Args: + sglang_tp_size (int): The number of tensor parallel workers. Defaults to 1. + sglang_pp_size (int): The number of pipeline parallel workers. Defaults to 1. + sglang_dp_size (int): The number of data parallel workers. Defaults to 1. + sglang_ep_size (int): The number of expert parallel workers. Defaults to 1. + sglang_enable_ep_moe (bool): Whether to enable expert parallelism for MoE. + Note: This argument has been removed in recent versions of SGLang. Defaults to False. + sglang_mem_fraction_static (Optional[float]): The fraction of GPU memory for the static allocation of model + weights and the KV cache memory pool. Try lowering this value if you encounter GPU out-of-memory errors. + Defaults to None. + sglang_context_length (Optional[int]): The maximum context length for the model. If None, the value from the + model's `config.json` will be used. Defaults to None. + sglang_disable_cuda_graph (bool): Disable CUDA graph for inference. Defaults to False. + sglang_quantization (Optional[str]): The quantization method to use. Defaults to None. + sglang_kv_cache_dtype (str): The data type for K/V cache storage. 'auto' will use the model's data type. + 'fp8_e5m2' and 'fp8_e4m3' are available for CUDA 11.8 and later. Defaults to 'auto'. + sglang_enable_dp_attention (bool): Enables data parallelism for the attention mechanism and tensor parallelism + for the feed-forward network (FFN). The data parallel size (dp_size) must equal the tensor parallel size + (tp_size). Currently supported for DeepSeek-V2/3 and Qwen2/3 MoE models. Defaults to False. + sglang_disable_custom_all_reduce (bool): Disable the custom all-reduce kernel and fall back to NCCL. Enabled by + default (True) for stability. Defaults to True. + sglang_speculative_algorithm (Optional[str]): The speculative decoding algorithm. Options include "EAGLE", + "EAGLE3", "NEXTN", "STANDALONE", "NGRAM". Defaults to None. + sglang_speculative_num_steps (Optional[int]): The number of steps to sample from the draft model during + speculative decoding. Defaults to None. + sglang_speculative_eagle_topk (Optional[int]): The number of tokens to sample from the draft model at each step + for the EAGLE2 algorithm. Defaults to None. + sglang_speculative_num_draft_tokens (Optional[int]): The number of tokens to sample from the draft model during + speculative decoding. Defaults to None. + """ + sglang_tp_size: int = 1 + sglang_pp_size: int = 1 + sglang_dp_size: int = 1 + sglang_ep_size: int = 1 + sglang_enable_ep_moe: bool = False + sglang_mem_fraction_static: Optional[float] = None + sglang_context_length: Optional[int] = None + sglang_disable_cuda_graph: bool = False + sglang_quantization: Optional[str] = None + sglang_kv_cache_dtype: str = 'auto' + sglang_enable_dp_attention: bool = False + sglang_disable_custom_all_reduce: bool = True + # speculative decoding + # e.g. EAGLE, EAGLE3, NEXTN + sglang_speculative_algorithm: Optional[str] = None + sglang_speculative_num_steps: Optional[int] = None + sglang_speculative_eagle_topk: Optional[int] = None + sglang_speculative_num_draft_tokens: Optional[int] = None + + def get_sglang_engine_kwargs(self): + kwargs = { + 'tp_size': self.sglang_tp_size, + 'pp_size': self.sglang_pp_size, + 'dp_size': self.sglang_dp_size, + 'ep_size': self.sglang_ep_size, + 'enable_ep_moe': self.sglang_enable_ep_moe, + 'mem_fraction_static': self.sglang_mem_fraction_static, + 'context_length': self.sglang_context_length, + 'disable_cuda_graph': self.sglang_disable_cuda_graph, + 'quantization': self.sglang_quantization, + 'kv_cache_dtype': self.sglang_kv_cache_dtype, + 'enable_dp_attention': self.sglang_enable_dp_attention, + 'disable_custom_all_reduce': self.sglang_disable_custom_all_reduce, + 'speculative_algorithm': self.sglang_speculative_algorithm, + 'speculative_num_steps': self.sglang_speculative_num_steps, + 'speculative_eagle_topk': self.sglang_speculative_eagle_topk, + 'speculative_num_draft_tokens': self.sglang_speculative_num_draft_tokens, + } + if self.task_type == 'embedding': + kwargs['task_type'] = 'embedding' + return kwargs + + +@dataclass +class InferArguments(MergeArguments, LmdeployArguments, SglangArguments, VllmArguments, BaseArguments): + """Arguments for model inference. + + A dataclass that extends BaseArguments, MergeArguments, VllmArguments, and LmdeployArguments to define all + arguments required for model inference. + + Args: + infer_backend (Literal['transformers', 'vllm', 'sglang', 'lmdeploy']): The inference acceleration + backend to use. Defaults to 'transformers'. + result_path (Optional[str]): The path to store inference results in JSONL format. If the file already exists, + new results will be appended. If None, results are saved in the checkpoint directory (if available) or + './result'. The final path will be printed to the console. Defaults to None. + write_batch_size (int): The batch size for writing results to `result_path`. A value of -1 means no limit. + Defaults to 1000. + metric (Optional[str]): The metric to use for evaluating inference results. Supported values are 'acc' and + 'rouge'. If None, no evaluation is performed. Defaults to None. + max_batch_size (int): The maximum batch size for inference, effective only when `infer_backend` is + 'transformers'. A value of -1 means no limit. Defaults to 1. + val_dataset_sample (Optional[int]): The number of samples to use from the inference dataset. If None, the + entire dataset is used. Defaults to None. + reranker_use_activation (bool): Whether to apply a sigmoid activation to the scores during reranker inference. + Defaults to True. + """ + # `pt` is used for swift3.x shell script compatibility. + infer_backend: Literal['vllm', 'transformers', 'sglang', 'lmdeploy', 'pt'] = 'transformers' + + result_path: Optional[str] = None + write_batch_size: int = 1000 + metric: Literal['acc', 'rouge'] = None + # for transformers engine + max_batch_size: int = 1 + + # only for inference + val_dataset_sample: Optional[int] = None + + # for reranker + reranker_use_activation: bool = True + + def _get_result_path(self, folder_name: str) -> str: + result_dir = self.ckpt_dir or f'result/{self.model_suffix}' + os.makedirs(result_dir, exist_ok=True) + result_dir = to_abspath(os.path.join(result_dir, folder_name)) + os.makedirs(result_dir, exist_ok=True) + time = dt.datetime.now().strftime('%Y%m%d-%H%M%S') + return os.path.join(result_dir, f'{time}.jsonl') + + def _init_result_path(self, folder_name: str) -> None: + if self.result_path is not None: + self.result_path = to_abspath(self.result_path) + return + # By default, a result_path file is automatically created + # when a validation or evaluation dataset is present. + if self._val_dataset_exists or getattr(self, 'eval_dataset', None): + self.result_path = self._get_result_path(folder_name) + logger.info(f'args.result_path: {self.result_path}') + + def _init_stream(self): + self.eval_human = not self._val_dataset_exists + logger.info(f'Setting args.eval_human: {self.eval_human}') + if self.stream is None: + self.stream = self.eval_human + if self.stream and self.num_beams != 1: + self.stream = False + logger.info('Setting args.stream: False') + + def _init_ddp(self): + if not is_dist(): + return + eval_human = getattr(self, 'eval_human', False) + assert not eval_human and not self.stream, ( + 'In DDP scenarios, interactive interfaces and streaming output are not supported.' + f'args.eval_human: {eval_human}, args.stream: {self.stream}') + self._init_device() + init_process_group(backend=self.ddp_backend, timeout=self.ddp_timeout) + + def __post_init__(self) -> None: + if self.infer_backend == 'pt': + self.infer_backend = 'transformers' # compat swift3.x + logger.warning('args.infer_backend: `pt` is deprecated, please use args.infer_backend: `transformers`.') + BaseArguments.__post_init__(self) + VllmArguments.__post_init__(self) + self._init_vllm_async_engine() + # Default to False for swift infer (non-encode tasks) + if self.vllm_use_async_engine is None: + self.vllm_use_async_engine = False + self._init_result_path('infer_result') + self._init_ddp() + + def _init_vllm_async_engine(self): + """Initialize vllm_use_async_engine based on task_type. + + Encode tasks (embedding, seq_cls, reranker, generative_reranker) require + async engine because vLLM's synchronous LLMEngine does not have the `encode` method. + + Note: This method only handles encode tasks. For non-encode tasks, the default value + should be set by subclasses (DeployArguments sets True, RolloutArguments uses + _set_default_engine_type, InferArguments defaults to False). + """ + # Task types that require vLLM's encode() method, which is only available in AsyncLLMEngine + encode_task_types = ('embedding', 'seq_cls', 'reranker', 'generative_reranker') + is_vllm_encode_task = self.infer_backend == 'vllm' and self.task_type in encode_task_types + + if is_vllm_encode_task: + if self.vllm_use_async_engine is None: + self.vllm_use_async_engine = True + elif not self.vllm_use_async_engine: + raise ValueError( + f'task_type={self.task_type} requires vllm_use_async_engine=True. ' + f'The synchronous vLLM LLMEngine does not support the `encode` method for encode tasks. ' + f'Please set --vllm_use_async_engine true or remove the explicit false setting.') diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/merge_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/merge_args.py new file mode 100644 index 0000000000000000000000000000000000000000..9b10d49e9e3cbca586b06fdbeddc0c457c17faf7 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/merge_args.py @@ -0,0 +1,23 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import dataclass + +from swift.utils import get_logger + +logger = get_logger() + + +@dataclass +class MergeArguments: + """A dataclass that holds configuration for merging models. + + This dataclass stores all the arguments needed to configure the model merging process. + + Args: + merge_lora (bool): Whether to merge LoRA adapters. This parameter supports `lora`, `llamapro`, and `longlora`. + Defaults to False. + safe_serialization (bool): Whether to use safetensors for serialization. Defaults to True. + max_shard_size (str): The maximum size of a single saved shard file. Defaults to '5GB'. + """ + merge_lora: bool = False + safe_serialization: bool = True + max_shard_size: str = '5GB' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/pretrain_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/pretrain_args.py new file mode 100644 index 0000000000000000000000000000000000000000..6d096511b066c4722148f46345547b46dd47e800 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/pretrain_args.py @@ -0,0 +1,10 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import dataclass + +from .sft_args import SftArguments + + +@dataclass +class PretrainArguments(SftArguments): + use_chat_template: bool = False + loss_scale: str = 'all' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/rlhf_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/rlhf_args.py new file mode 100644 index 0000000000000000000000000000000000000000..8b1097c35345cf87189e45b8bf9d377398efb07a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/rlhf_args.py @@ -0,0 +1,772 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +from dataclasses import dataclass, field +from typing import Any, Dict, List, Literal, Optional + +from swift.model import MODEL_MAPPING +from swift.rlhf_trainers import GRPOArgumentsMixin +from swift.template import TEMPLATE_MAPPING +from swift.utils import get_current_device, get_logger, is_master, is_mp, json_parse_to_dict, set_default_ddp_config +from .sft_args import SftArguments + +logger = get_logger() +rlhf_support_vllm_types = ['grpo', 'gkd'] + + +@dataclass +class RewardModelArguments: + """Arguments pertaining to the reward model. + + Args: + reward_model (Optional[List[str]]): The model ID or a local path to the reward model. Same as the `model` + argument. Defaults to None. + reward_adapters (List[str]): The path(s) to LoRA adapter weights to be loaded for the reward model. Useful for + using LoRA weights from SFT as the reward model. Defaults to an empty list (`[]`). + reward_model_type (Optional[List[str]]): The model type of the reward model. Same as the `model_type` argument. + If not specified, it's often inferred. Defaults to None. + reward_model_revision (Optional[List[str]]): The specific model version to use for the reward model. Same as + the `model_revision` argument. Defaults to None. + reward_template (Optional[List[str]]): The template to use for the reward model. Defaults to None. + """ + reward_model: Optional[List[str]] = None + reward_adapters: List[str] = field(default_factory=list) + reward_model_type: Optional[List[str]] = field( + default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'}) + reward_model_revision: Optional[List[str]] = None + reward_template: Optional[List[str]] = field( + default=None, metadata={'help': f'template choices: {list(TEMPLATE_MAPPING.keys())}'}) + + +@dataclass +class TeacherModelArguments: + """Arguments for configuring the teacher model. + + Args: + teacher_model (Optional[str]): The model ID or a local path to the teacher model. Analogous to the main + `model` argument. For GKD, there are three modes: + - Not set (None): Self-distillation with dynamic teacher (teacher = current student weights). + - Same as `model` with LoRA training: Self-distillation with fixed teacher. Automatically optimized + to use `disable_adapter()` to get base model logits without loading an extra model. + - Different from `model`: Standard GKD with an independent frozen teacher model. + Defaults to None. + teacher_adapters (List[str]): A list of paths to LoRA weights. These weights, often produced by SFT, are loaded + to form the teacher model. Defaults to an empty list (`[]`). + teacher_model_type (Optional[str]): The model type of the teacher model. If not specified, it's often inferred. + Analogous to the main `model_type` argument. Defaults to None. + teacher_model_revision (Optional[str]): The specific model version of the teacher model to use. Analogous to + the main `model_revision` argument. Defaults to None. + teacher_deepspeed (Optional[str]): The teacher model's deepspeed configuration. This can be a JSON file path or + one of the following values: 'zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload'. If not + provided, it defaults to using the same DeepSpeed configuration as the main training model. Analogous to + the main `deepspeed` argument. + teacher_model_server (Optional[str]): The URL of the teacher model server (e.g., 'http://localhost:8000'). + When set, the teacher logprobs will be fetched from the external API service (e.g., swift deploy, vLLM) + instead of loading a local teacher model. This enables using larger teacher models or services hosted + remotely. When this is set, `teacher_model` is not required. Defaults to None. + offload_teacher_model (bool): Whether to offload the teacher model to CPU memory to save VRAM during GKD + or OPD-RL training. When enabled, the teacher model is loaded to GPU only during forward pass and + offloaded back to CPU afterwards. Defaults to False. + """ + teacher_model: Optional[str] = None + teacher_adapters: List[str] = field(default_factory=list) + teacher_model_type: Optional[str] = field( + default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'}) + teacher_model_revision: Optional[str] = None + teacher_deepspeed: Optional[str] = field( + default=None, + metadata={ + 'help': + 'DeepSpeed configuration for teacher model. ' + 'Can be a path to a json file or one of: zero0, zero1, zero2, zero3, zero2_offload, zero3_offload' + }) + teacher_model_server: Optional[str] = field( + default=None, + metadata={ + 'help': + 'URL of the teacher model server (e.g., http://localhost:8000). ' + 'When set, teacher logprobs are fetched via API instead of loading a local model. ' + 'Supports multi-teacher via JSON list of {url, tags}.' + }) + offload_teacher_model: bool = False + + +@dataclass +class PPOArguments: + """Arguments for configuring the PPO training. + + Args: + num_ppo_epochs (int): Number of epochs to train. Defaults to 4. + whiten_rewards (bool): Whether to whiten the rewards. Defaults to False. + kl_coef (float): KL coefficient. Defaults to 0.05. + cliprange (float): Clip range. Defaults to 0.2. + vf_coef (float): Value function coefficient. Defaults to 0.1. + cliprange_value (float): Clip range for the value function. Defaults to 0.2. + gamma (float): Discount factor. Defaults to 1.0. + lam (float): Lambda value for GAE. Defaults to 0.95. + num_mini_batches (int): Defaults to 1. + local_rollout_forward_batch_size (int): Defaults to 64. + num_sample_generations (int): Number of generations. Defaults to 10. + response_length (Optional[int]): (Deprecated) Compatibility parameter. Use `max_completion_length` instead. + Defaults to None. + missing_eos_penalty (Optional[float]): Defaults to None. + """ + num_ppo_epochs: int = 4 + whiten_rewards: bool = False + kl_coef: float = 0.05 + cliprange: float = 0.2 + vf_coef: float = 0.1 + cliprange_value: float = 0.2 + gamma: float = 1.0 + lam: float = 0.95 + + num_mini_batches: int = 1 + local_rollout_forward_batch_size: int = 64 + num_sample_generations: int = 10 + response_length: Optional[int] = None # compat. use max_completion_length instead + missing_eos_penalty: Optional[float] = None + + +@dataclass +class GRPOArguments(GRPOArgumentsMixin): + """A dataclass for configuring GRPO training. + + These arguments control the hyperparameters specific to the GRPO algorithm. + + Args: + num_generations (int): The number of completions to generate for each prompt. This corresponds to the G value + in the GRPO paper. The total generation batch size (e.g., `generation_batch_size` or `steps_per_generation + * per_device_batch_size * num_processes`) must be divisible by this number. Defaults to 8. + num_generations_eval (Optional[int]): Number of generations to sample during evaluation. This allows + using fewer generations during evaluation to save computation. If `None`, uses the value of + `num_generations`. Defaults to None. + reward_funcs (List[str]): A list of reward function names to use for the GRPO algorithm. Available built-in + options include 'accuracy', 'format', 'cosine', 'repetition', and 'soft_overlong' + (see swift/rewards/orm.py). Custom reward functions can also be defined. Defaults to an empty list. + reward_weights (List[float]): A list of weights for each reward source. The length must match the total number + of reward functions (from `reward_funcs`) plus any external reward models. If `None`, all rewards are + weighted equally with a value of 1.0. Note: If an external `--reward_model` is used, it is treated as the + last reward source in the sequence. Defaults to None. + log_completions (bool): Whether to log the model's generated completions during training. This is designed to + be used with an experiment tracker like WandB or SwanLab (`--report_to wandb`/`swanlab`). If enabled + without a tracker, completions are saved to `completions.jsonl` in the checkpoint directory. Defaults to + False. + num_iterations (int): The number of update steps to perform for each data sample. This corresponds to the K + value in the GRPO paper. Defaults to 1. + truncation_strategy (Literal['delete', 'left', 'right', 'split', None]): The strategy for handling input + sequences that exceed `max_length`. Supported options: 'delete' to discard the sample, 'left' to truncate + from the beginning, 'right' to truncate from the end. Defaults to None, and then sets to 'left' in the + `_init_grpo` function. + Note that for multimodal models, left pruning may prune multimodal tokens, causing shape mismatch errors + in the forward feed. Using the `delete` method will resample other data from the original dataset to + supplement excessively long data and examples with encoding failures. + """ + num_generations: int = 8 # G in the GRPO paper + reward_funcs: List[str] = field(default_factory=list) + reward_weights: List[float] = None + log_completions: bool = False + + # multi step + num_iterations: int = 1 + + truncation_strategy: Literal['delete', 'left', 'right', 'split', None] = None + + +@dataclass +class RLHFArguments(TeacherModelArguments, GRPOArguments, PPOArguments, RewardModelArguments, SftArguments): + """A dataclass holding arguments for Reinforcement Learning from Human Feedback. + + Args: + rlhf_type (str): The type of human alignment algorithm to use. Supports 'dpo', 'orpo', 'simpo', 'kto', 'cpo', + 'rm', 'ppo', 'grpo', and 'gkd'. Defaults to 'dpo'. + ref_model (Optional[str]): The model path for the reference model. Required when using 'dpo', 'kto', 'ppo', + or 'grpo' with full-parameter training. Defaults to None, which will set it to the value of the `--model` + argument. + ref_adapters (List[str]): LoRA adapters for the reference model. If you are using LoRA weights from SFT for + DPO/KTO/GRPO, set both `--adapters` and `--ref_adapters` to the SFT checkpoint path. When resuming from an + RLHF checkpoint, set `--resume_from_checkpoint` to the RLHF checkpoint and `--ref_adapters` to the SFT + checkpoint. Defaults to an empty list. + ref_model_type (Optional[str]): The model type of the reference model. Same as `model_type`. Defaults to None. + ref_model_revision (Optional[str]): The model revision of the reference model. Same as `model_revision`. + Defaults to None. + beta (Optional[float]): The beta parameter for RLHF, controlling the deviation from the reference model. + A higher value implies less deviation. If None, uses algorithm-specific defaults: 2.0 for 'simpo', 0.04 + for 'grpo', 0.5 for 'gkd', and 0.1 for others. Defaults to None. + label_smoothing (float): The label smoothing value for DPO. A value of 0 disables it. Defaults to 0. + max_completion_length (int): The maximum generation length for GRPO/PPO/GKD algorithms. Defaults to 512. + loss_scale (Optional[str]): Overrides the template parameter. During RLHF training, this defaults to + 'last_round'. + rpo_alpha (Optional[float]): The alpha parameter from the RPO paper, controlling the weight of the SFT loss + (NLL term). The loss is calculated as `dpo_loss + rpo_alpha * sft_loss`. If None, the SFT loss is not + included. + ld_alpha (Optional[float]): The alpha parameter from the LD-DPO paper, which weights the log probabilities of + the sequence part beyond the common prefix to mitigate length preference. Defaults to None. + discopop_tau (float): The temperature parameter from the DiscoPOP paper, used to scale the log-ratio. Effective + when `loss_type` is 'discopop'. Defaults to 0.05. + loss_type (Optional[List[str]]): The type of loss function. Defaults to algorithm-specific values (e.g., + 'sigmoid' for DPO). Multiple values can be passed for mixed training (MPO), which requires `loss_weights` + to be set. + loss_weights (Optional[List[float]]): When multiple `loss_type` values are set for DPO, this specifies the + weights for each loss term. Defaults to None. + cpo_alpha (float): The coefficient for the NLL loss in the CPO/SimPO loss function. Defaults to 1.0. + simpo_gamma (float): The reward margin term in the SimPO algorithm. The paper suggests a value between 0.5 and + 1.5. Defaults to 1.0. + desirable_weight (float): In KTO, the weight applied to the desirable loss to counteract data imbalance. + Defaults to 1.0. + undesirable_weight (float): In KTO, the weight applied to the undesirable loss to counteract data imbalance. + Defaults to 1.0. + temperature (float): The temperature for sampling, used in PPO, GRPO, and GKD algorithms. Defaults to 0.9. + center_rewards_coefficient (Optional[float]): Used for Reward Model (RM) training. A coefficient to encourage + the reward model to output rewards with a mean of zero. A value of 0.01 is recommended. Defaults to None. + sft_alpha (float): The weight for the SFT loss component in GKD. The final loss is calculated as + gkd_loss + sft_alpha * sft_loss`. Defaults to 0. + lmbda (float): The lambda parameter for GKD, balancing policy and value losses. Defaults to 0.5. + seq_kd (bool): Deprecated. Sequential KD (teacher-generated responses) is not implemented. + gkd_logits_topk (Optional[int]): The number of top-k logits to use for KL divergence computation in GKD. + If None, uses full vocabulary for KL computation (more accurate but memory-intensive). + If set to a positive integer, only top-k teacher logits are used (more efficient). + When using `teacher_model_server`, this is limited by the server's `max_logprobs` setting + (vLLM default is 20, can be increased with `--max-logprobs`). Defaults to None. + max_new_tokens (Optional[int]): A backward-compatibility argument. Please use `max_completion_length` instead. + Defaults to None. + """ + rlhf_type: Literal['dpo', 'orpo', 'simpo', 'kto', 'cpo', 'rm', 'ppo', 'grpo', 'gkd'] = 'dpo' + ref_model: Optional[str] = None + ref_adapters: List[str] = field(default_factory=list) + ref_model_type: Optional[str] = field( + default=None, metadata={'help': f'model_type choices: {list(MODEL_MAPPING.keys())}'}) + ref_model_revision: Optional[str] = None + + beta: Optional[float] = None + label_smoothing: float = 0 + max_completion_length: int = 512 + loss_scale: Optional[str] = None # 'last_round' + # DPO + rpo_alpha: Optional[float] = None + ld_alpha: Optional[float] = None # α parameter from the LD-DPO paper + discopop_tau: float = 0.05 # τ/temperature parameter from the DiscoPOP paper + loss_type: Optional[List[str]] = None + loss_weights: Optional[List[float]] = None + # CPO + cpo_alpha: float = 1. + # SimPO + simpo_gamma: float = 1 + # KTO + desirable_weight: float = 1.0 + undesirable_weight: float = 1.0 + # PPO/GRPO/GKD + temperature: float = 0.9 + # RM + center_rewards_coefficient: Optional[float] = None + # GKD + sft_alpha: float = 0 + lmbda: float = 0.5 + seq_kd: bool = False # Deprecated + gkd_logits_topk: Optional[int] = None + # compat + max_new_tokens: Optional[int] = None # use max_completion_length instead + + def _prepare_training_args(self, training_args: Dict[str, Any]) -> None: + if self.rlhf_type == 'ppo': + training_args['world_size'] = self.global_world_size + + def __post_init__(self): + self._process_loss_type() + self._init_grpo() + self._init_rm() + self._init_simpo() + self._init_max_completion_length() + self._init_padding_side() + self._set_default() + self._init_rollout() + self._init_teacher_deepspeed() + GRPOArguments.__post_init__(self) + SftArguments.__post_init__(self) + self._check_sequence_parallel() + self._check_teacher() + self._check_grpo() + self._check_gkd() + + if isinstance(self.ref_adapters, str): + self.ref_adapters = [self.ref_adapters] + if self.rlhf_type == 'grpo' and self.beta == 0.0: + self.ref_model = None + elif self.rlhf_type in ['dpo', 'kto', 'ppo', 'grpo'] and self.tuner_type == 'full': + self.ref_model = self.ref_model or self.model + self.ref_model_type = self.ref_model_type or self.model_type + self.ref_model_revision = self.ref_model_revision or self.model_revision + elif self.ref_model is not None: + raise ValueError('CPO/ORPO or LoRA training does not require a ref_model to be passed in.') + + def _set_loss_scale(self): + if self.loss_scale is None: + if self.rlhf_type == 'orpo' and not self.model_meta.is_multimodal: + # Avoid padding labels during the model's forward pass in multimodal models. + # Some multimodal models do not expand the image pad token. + self.loss_scale = 'default' + elif self.rlhf_type in ('grpo', 'gkd'): + if self.multi_turn_scheduler: + self.loss_scale = 'default' + else: + self.loss_scale = 'last_round' + else: + self.loss_scale = 'last_round' + super()._set_loss_scale() + + def _process_loss_type(self): + if self.loss_type is None: + return + + if isinstance(self.loss_type, list): + num_loss_types = len(self.loss_type) + if num_loss_types > 1: + assert self.rlhf_type == 'dpo', (f'Multiple loss types ({self.loss_type}) are only supported for DPO. ' + f'Current rlhf_type: {self.rlhf_type}.') + from trl.trainer.dpo_config import DPOConfig + assert 'loss_weights' in DPOConfig.__dict__, ( + 'Multiple loss types requires trl >= 0.20, please install trl `pip install -U trl`') + + if hasattr(self.loss_type, '__len__') and len(self.loss_type) == 1: + self.loss_type = self.loss_type[0] + + # Validate loss_type + if self.loss_weights is not None: + assert self.rlhf_type == 'dpo' + loss_types = self.loss_type if isinstance(self.loss_type, list) else [self.loss_type] + if len(self.loss_weights) != len(loss_types): + raise ValueError(f'Length of loss_weights list ({self.loss_weights}) must match number of loss types ' + f'({loss_types}).') + + def _init_grpo(self): + if self.rlhf_type != 'grpo': + return + if self.cached_dataset or self.cached_val_dataset: + raise ValueError('cached_dataset is not supported for GRPO.') + if self.use_vllm: + set_default_ddp_config() + if self.async_generate or not self.use_vllm or self.vllm_mode == 'server': + self.sleep_level = 0 + self.remove_unused_columns = False + logger.info(f'Setting args.remove_unused_columns: {self.remove_unused_columns}') + if self.truncation_strategy is None: + self.truncation_strategy = 'left' + if self.truncation_strategy not in {'left', 'delete'}: + raise ValueError("GRPO requires `truncation_strategy 'left' or 'delete'`, " + f"Current value: `truncation_strategy='{self.truncation_strategy}'`.") + if self.beta is None: + self.beta = 0.04 # https://arxiv.org/abs/2402.03300 + if self.async_generate: + logger.info('Using async mode. This is a approximate version which ' + 'will use the old weights to generate responses to accelerate. ' + 'This will ignore the `CLIP` of advantages, if you found the training ' + 'is unstable, you may consider using --async_generate false.') + if 'soft_overlong' in self.reward_funcs: + assert self.soft_cache_length is not None, \ + 'The soft_cache_length must be set when using soft overlong rewards.' + if self.soft_max_length is None: + self.soft_max_length = self.max_completion_length + logger.info(f'Auto-configured soft_max_length = max_completion_length {self.max_completion_length}') + + if self.kl_in_reward is None: + if self.advantage_estimator == 'grpo': + self.kl_in_reward = False + elif self.advantage_estimator in ['rloo', 'reinforce_plus_plus']: + self.kl_in_reward = True + else: + raise ValueError(f'Invalid advantage_estimator: {self.advantage_estimator}') + + # disable normalization, REAL https://arxiv.org/abs/2602.05630 + if self.loss_type == 'real': + self.scale_rewards = 'none' + logger.warning( + f"[REAL] scale_rewards='{self.scale_rewards}' is ignored. " + "It will be forced to 'none' because 'loss_type = real' does not support reward normalization.") + + if self.scale_rewards is None: + if self.advantage_estimator == 'grpo': + self.scale_rewards = 'group' + elif self.advantage_estimator == 'rloo': + self.scale_rewards = 'none' + elif self.advantage_estimator == 'reinforce_plus_plus': + self.scale_rewards = 'batch' + else: + raise ValueError(f'Invalid advantage_estimator: {self.advantage_estimator}') + + def _check_teacher(self): + self._teacher_use_disable_adapter = False + + if self.rlhf_type not in ['grpo', 'gkd']: + if self.teacher_model is not None or self.teacher_model_server is not None: + logger.warning(f'teacher_model / teacher_model_server is ignored for rlhf_type={self.rlhf_type!r} ' + '(only used by gkd and grpo/OPD-RL).') + return + teacher_set = self.teacher_model is not None or self.teacher_model_server is not None + if not teacher_set: + if self.rlhf_type == 'gkd': + logger.info('No teacher_model specified. Using self-distillation mode (teacher = student).') + if self.use_liger_kernel: + raise ValueError('Self-distillation mode with liger kernel loss is not supported yet') + if self.rlhf_type == 'grpo' and self.num_generations == 1: + raise ValueError('num_generations must be greater than 1 for GRPO') + return + + if self.rlhf_type == 'grpo' and self.use_liger_kernel: + raise ValueError('OPD-RL is not supported with use_liger_kernel.') + + if self.teacher_model is not None and self.teacher_model_server is not None: + raise ValueError('setting both `teacher_model` and `teacher_model_server` is not supported.') + + # Validate teacher_model_server: accept single URL or JSON multi-teacher config. + if self.teacher_model_server is not None: + from swift.rlhf_trainers.gkd_helpers import parse_teacher_model_server + + # Parse early to fail fast on invalid JSON; result is re-parsed by the trainer. + parse_teacher_model_server(self.teacher_model_server) + + # Self-distillation: teacher_model == student model + if self.teacher_model is not None and self.teacher_model == self.model: + if self.tuner_type == 'lora': + logger.info('LoRA + same teacher_model: using disable_adapter() for fixed teacher (no extra model).') + self._teacher_use_disable_adapter = True + self.teacher_model = None + else: + # Full training + same teacher_model: a separate frozen copy will be loaded as fixed teacher. + pass + + def _init_rollout(self): + if self.rlhf_type not in rlhf_support_vllm_types: + return + + if self.use_vllm and os.getenv('SWIFT_AUDIO_LOAD_BACKEND') is None: + # align with vLLM audio load backend + os.environ['SWIFT_AUDIO_LOAD_BACKEND'] = 'soundfile_pyav' + + if self.vllm_mode is not None and not self.use_vllm: + raise ValueError('vllm_mode is not supported when use_vllm is false') + + if self.vllm_mode is None and self.use_vllm: + raise ValueError('vllm_mode is required when use_vllm is true') + + self._init_external_vllm() + + if self.vllm_mode == 'server': + assert not self.use_vllm or self.vllm_server_host is not None or self.vllm_server_base_url is not None + + if self.async_generate: + assert self.vllm_mode == 'server', 'async generate require vllm_mode == server, ' + 'please deploy vLLM server by `swift rollout` and assign with `vllm_server_host` ' + 'for more infomations, please check ' + 'https://swift.readthedocs.io/en/latest/Instruction/GRPO/getstarted/GRPO.html' + + if not self.use_vllm and self.vllm_tensor_parallel_size != 1: + self.vllm_tensor_parallel_size = 1 + logger.warning('set vllm_tensor_parallel_size to 1 since use_vllm false') + self._external_vllm_warning() + + def _init_padding_side(self): + if self.rlhf_type in {'ppo', 'gkd'}: + self.padding_side = 'left' + # TODO: streaming, MLLM + + def _init_max_completion_length(self): + max_completion_length = self.response_length or self.max_new_tokens or self.max_completion_length + self.max_completion_length = self.max_new_tokens = self.response_length = max_completion_length + + def _init_metric_for_best_model(self): + if self.rlhf_type == 'grpo' and self.metric_for_best_model is None: + self.metric_for_best_model = 'reward' + super()._init_metric_for_best_model() + if self.rlhf_type == 'ppo': + self.metric_for_best_model = None + self.greater_is_better = None + + def _init_simpo(self): + if self.rlhf_type != 'simpo': + return + + self.rlhf_type = 'cpo' + if self.loss_type is None: + self.loss_type = 'simpo' + if self.beta is None: + self.beta = 2. + + def _init_rm(self): + if self.rlhf_type == 'rm': + self.task_type = 'seq_cls' + self.num_labels = 1 + + def _init_external_vllm(self): + if self.rlhf_type not in rlhf_support_vllm_types or (self.vllm_server_host is None + and self.vllm_server_base_url is None): + return + from swift.rlhf_trainers import VLLMClient + if is_master(): + logger.info('Start connecting to vLLM server') + self.vllm_client = VLLMClient( + base_urls=self.vllm_server_base_url, + hosts=self.vllm_server_host, + server_ports=self.vllm_server_port, + group_ports=self.vllm_server_group_port, + connection_timeout=self.vllm_server_timeout) + self.vllm_client.close_communicator() + self.vllm_client.init_communicator(device=get_current_device()) + logger.info('Connected to vLLM server') + + def _set_default(self): + if self.beta is None: + if self.rlhf_type == 'gkd': + self.beta = 0.5 + else: + self.beta = 0.1 + if self.loss_type is None: + if self.rlhf_type in ['dpo', 'cpo']: + self.loss_type = 'sigmoid' # else None + elif self.rlhf_type in ['kto']: + self.loss_type = 'kto' + elif self.rlhf_type == 'grpo': + self.loss_type = 'grpo' + if self.gradient_accumulation_steps is None: + if self.rlhf_type == 'grpo': + self.gradient_accumulation_steps = 1 + logger.info('Setting default gradient_accumulation_steps to 1 for GRPO.') + + def _check_grpo(self): + if self.rlhf_type != 'grpo': + return + import importlib.metadata + import trl + from packaging import version + trl_version = version.parse(trl.__version__) + assert trl_version >= version.parse('0.20'), ('Your current version of `trl` is outdated. ' + 'Please update it by running: pip install -U trl') + if is_mp() and self.use_vllm: + raise ValueError('GRPO with vLLM is not compatible with `device_map`. ' + 'Please set NPROC_PER_NODE equal to num_processes.') + if self.use_liger_kernel: + liger_kernel_version = version.parse(importlib.metadata.version('liger-kernel')) + if liger_kernel_version < version.parse('0.7.0'): + raise ValueError('Please update liger-kernel to 0.7.0 or later: pip install -U liger-kernel') + if self.delta is not None: + raise ValueError('Liger loss does not support two-sided GRPO loss yet.') + if self.sequence_parallel_size > 1: + raise ValueError('Liger loss does not support sequence parallel yet.') + if self.padding_free: + raise ValueError('Liger loss does not support padding free yet.') + if self.top_entropy_quantile < 1.0: + raise ValueError('Liger loss does not support entropy mask yet.') + if self.log_entropy: + raise ValueError('Liger loss does not support log entropy yet.') + if self.off_policy_sequence_mask_delta is not None: + raise ValueError('Liger loss does not support off-policy sequence masking yet.') + assert self.importance_sampling_level in [ + 'token', 'sequence' + ], ('Liger loss currently only support token-level and sequence-level importance sampling. ' + 'Please set `importance_sampling_level` to `token` or `sequence`.') + if self.advantage_estimator != 'grpo': + raise ValueError('Liger loss currently only support grpo advantage estimator') + + if self.async_generate and self.multi_turn_scheduler is not None: + raise NotImplementedError('Currently, async_generate is not supported with multi-turn functionality.') + + self._check_opd_rl() + self._check_rlsd() + self._check_sdar() + + def _check_rlsd(self): + """Validate RLSD (Self-Distilled RLVR) advantage reweighting parameters. + + RLSD reweights the per-sequence GRPO advantage per token using the teacher-vs-student + logprob gap. It relies on an OPSD teacher forward driven by a per-sample ``teacher_prompt`` + column (privileged context, e.g. question + reference solution / GT answer). It also needs + an environment reward (``reward_funcs``) because the reweight direction uses ``sign(A)``. + + Two teacher modes are supported (mirroring the reference's ``freeze_teacher_model`` flag): + + * **Dynamic self-distillation** (no ``--teacher_model``, default): teacher = current policy + evaluated on the privileged prompt (paper's same-model, two-contexts Δ_t; equivalent to + the reference's ``freeze_teacher_model=false``). + * **Frozen local teacher** (``--teacher_model ``): a separate model scores the + privileged prompt under ``no_grad`` and is never in the optimizer, so it stays frozen + throughout training. Point ``--teacher_model`` at the initial policy checkpoint to + reproduce the reference's ``freeze_teacher_model=true`` (frozen θ₀ anchor). + """ + if self.advantage_reweight != 'rlsd': + return + if not (0.0 <= self.rlsd_lambda <= 1.0): + raise ValueError(f'rlsd_lambda must be in [0, 1], got {self.rlsd_lambda}.') + if self.rlsd_reweight_clip_range < 0: + raise ValueError(f'rlsd_reweight_clip_range (eps_w) must be >= 0, got {self.rlsd_reweight_clip_range}.') + if not self.reward_funcs: + raise ValueError('advantage_reweight=rlsd requires reward_funcs: the reweight direction uses ' + 'sign(A) of the GRPO advantage, which needs an environment reward.') + if self.use_liger_kernel: + raise ValueError('advantage_reweight=rlsd is not compatible with use_liger_kernel ' + '(the fused loss path bypasses the per-token RLSD reweight).') + if self.loss_type in ['real', 'fipo']: + raise ValueError(f'advantage_reweight=rlsd does not support loss_type={self.loss_type!r} ' + '(it reduces the advantage to a per-sequence scalar).') + if self.off_policy_sequence_mask_delta is not None: + raise ValueError('advantage_reweight=rlsd does not support off_policy_sequence_mask_delta.') + # Local frozen teacher (--teacher_model) is allowed: it drives the OPSD teacher forward on + # the privileged prompt via the same code path as OPD-RL, and use_rlsd bypasses OPD-RL's + # additive term (grpo_trainer.py:266) so there is no double-count. The API-teacher path + # (--teacher_model_server) scores tokens over a remote inference call whose tokenizer + # alignment with the student is not guaranteed, so keep it gated for RLSD's per-token reweight. + if self.teacher_model_server is not None: + raise ValueError('advantage_reweight=rlsd does not support --teacher_model_server (remote API ' + 'teacher). Use --teacher_model for a frozen local teacher, or omit ' + 'both to use dynamic self-distillation (teacher = current policy).') + if self.beta not in (0, 0.0): + logger.warning(f'advantage_reweight=rlsd: the reference RLSD config disables KL (beta=0), ' + f'but beta={self.beta}. Consider setting --beta 0.') + if self.teacher_model is not None: + logger.info(f'advantage_reweight=rlsd: frozen-teacher mode — --teacher_model={self.teacher_model!r} scores ' + f'the `teacher_prompt` column under no_grad and stays frozen throughout training (matches the ' + f'RLSD reference freeze_teacher_model=true when this checkpoint is the initial policy).') + else: + logger.info('advantage_reweight=rlsd: dynamic self-distillation mode — expecting a `teacher_prompt` ' + 'column in the dataset; teacher = current policy on the privileged prompt.') + + def _check_sdar(self): + """Validate SDAR (Self-Distilled Agentic RL) confidence-gated distillation parameters. + + SDAR adds a confidence-gated teacher distillation auxiliary loss to the GRPO policy loss: + ``loss = policy_loss + sdar_loss_coef * token-mean(sigmoid(gate_beta*Δ) * Δ)`` with + ``Δ = logP_T - logP_S``. It reuses the OPSD teacher forward driven by a per-sample + ``teacher_prompt`` column (privileged / skill-augmented context). Unlike RLSD it does NOT + modify the advantage, and the reference keeps a small KL (so KL is not forced off). + + Two teacher modes are supported (mirroring the reference's self-distillation setup): + + * **Dynamic self-distillation** (no ``--teacher_model``, default): teacher = current policy + on the privileged ``teacher_prompt`` (the reference ALFWorld/Search/WebShop setting). + * **Frozen local teacher** (``--teacher_model ``): a separate model scores the + privileged prompt under ``no_grad`` and stays frozen throughout training. + """ + if self.sdar_loss_coef <= 0: + return + if self.sdar_gate_beta <= 0: + raise ValueError(f'sdar_gate_beta must be > 0, got {self.sdar_gate_beta}.') + if self.use_liger_kernel: + raise ValueError('sdar_loss_coef>0 is not compatible with use_liger_kernel (the fused Liger ' + 'loss path bypasses the per-token loss where the SDAR term is added).') + if self.advantage_reweight == 'rlsd': + raise ValueError('sdar_loss_coef>0 cannot be combined with advantage_reweight=rlsd: SDAR adds a ' + 'confidence-gated distillation loss while RLSD reweights the advantage — pick one ' + 'teacher-distillation mechanism.') + if self.teacher_model_server is not None: + raise ValueError('sdar_loss_coef>0 does not support --teacher_model_server (remote API teacher). ' + 'Use --teacher_model for a frozen local teacher, or omit both to use ' + 'dynamic self-distillation (teacher = current policy on the `teacher_prompt`).') + if self.teacher_model is not None: + logger.info(f'SDAR (sdar_loss_coef={self.sdar_loss_coef}): frozen-teacher mode — ' + f'--teacher_model={self.teacher_model!r} scores the `teacher_prompt` column under no_grad ' + f'and stays frozen throughout training.') + else: + logger.info(f'SDAR (sdar_loss_coef={self.sdar_loss_coef}): dynamic self-distillation mode — expecting ' + 'a `teacher_prompt` column in the dataset; teacher = current policy on the privileged ' + 'prompt. Without teacher logps the SDAR loss is skipped (behaves as plain GRPO).') + + def _check_opd_rl(self): + """Fail-fast OPD-RL (teacher distillation on GRPO) parameter compatibility. + + A teacher turns GRPO into OPD-RL, where the teacher signal is a *per-token* advantage + (the signed teacher log-ratio). Features that require a *per-sequence* advantage (typically + sign-based judgments) or reward variance are incompatible; reject them here rather than + deep inside the loss / advantage code. ``_check_teacher`` has already run, so + ``_teacher_use_disable_adapter`` is resolved. + """ + opd_rl = ( + self.teacher_model is not None or self.teacher_model_server is not None + or self._teacher_use_disable_adapter) + if not opd_rl: + return + # loss types / masks that reduce the advantage to a per-sequence scalar (sign-based). + if self.loss_type in ['real', 'fipo']: + raise ValueError(f'OPD-RL (teacher) does not support loss_type={self.loss_type!r} ' + '(it needs a per-sequence advantage).') + if self.off_policy_sequence_mask_delta is not None: + raise ValueError('OPD-RL (teacher) does not support off_policy_sequence_mask_delta ' + '(it needs a per-sequence advantage).') + # Pure distillation (no reward functions): the base GRPO advantage is 0, so reward-variance + # driven features have no signal to act on. + if not self.reward_funcs: + if self.dynamic_sample: + raise ValueError('dynamic_sample requires reward_funcs (it filters groups by reward std); ' + 'pure OPD-RL distillation has no reward variance.') + if self.scale_rewards == 'gdpo': + raise ValueError("scale_rewards='gdpo' requires reward_funcs; pure OPD-RL distillation has none.") + + def _external_vllm_warning(self): + if self.rlhf_type not in rlhf_support_vllm_types or not self.vllm_server_host: + return + + if self.vllm_max_model_len is not None: + logger.warning( + "Configuration conflict: 'vllm_max_model_len=%s' is ignored for external vLLM. " + 'Please specify it when launching the inference service: ' + '`swift rollout --vllm_max_model_len `', self.vllm_max_model_len) + + def _check_padding_free(self): + super()._check_padding_free() + if self.padding_free or self.packing: + supported_types = ['grpo', 'dpo', 'kto', 'gkd'] + if self.rlhf_type not in supported_types: + raise NotImplementedError( + f"The current rlhf_type '{self.rlhf_type}' does not support padding_free/packing. " + 'Please set --padding_free/packing to false.') + + def _check_sequence_parallel(self): + if self.sequence_parallel_size > 1: + supported_types = ['grpo', 'dpo'] + if self.rlhf_type not in supported_types: + raise NotImplementedError( + f"The current rlhf_type '{self.rlhf_type}' does not support sequence_parallel. " + 'Please set --sequence_parallel_size to 1.') + + def _init_teacher_deepspeed(self): + """Initialize teacher_deepspeed configuration similar to _init_deepspeed in SftArguments""" + if not self.teacher_deepspeed: + return + + # Get the same ds_config_folder as main model + ds_config_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'config')) + deepspeed_mapping = { + name: f'{name}.json' + for name in ['zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload'] + } + + # Check if teacher_deepspeed is a predefined name + for ds_name, ds_config in deepspeed_mapping.items(): + if self.teacher_deepspeed == ds_name: + self.teacher_deepspeed = os.path.join(ds_config_folder, ds_config) + break + + # Parse the config file to dict + self.teacher_deepspeed = json_parse_to_dict(self.teacher_deepspeed) + logger.info(f'Using teacher_deepspeed config: {self.teacher_deepspeed}') + + def _check_gkd(self): + if self.rlhf_type != 'gkd': + return + if is_mp() and self.use_vllm: + raise ValueError('GKD with vLLM is not compatible with `device_map`. ' + 'Please set NPROC_PER_NODE equal to num_processes.') + + if self.async_generate: + raise NotImplementedError('Currently, async_generate is not supported for GKD.') + + # seq_kd (teacher-generated responses) is not implemented; raise early. + if self.seq_kd: + raise NotImplementedError('seq_kd=True (Sequential KD with teacher generation) is deprecated.') + + # When using teacher_model_server, gkd_logits_topk is required (API only returns top-k logprobs) + if self.teacher_model_server is not None: + if self.gkd_logits_topk is None: + raise ValueError('gkd_logits_topk is required when using teacher_model_server') + + # Validate gkd_logits_topk + if self.gkd_logits_topk is not None and self.gkd_logits_topk <= 0: + raise ValueError(f'gkd_logits_topk must be a positive integer, got {self.gkd_logits_topk}') + + if self.gkd_logits_topk is not None and self.use_liger_kernel: + raise ValueError('gkd_logits_topk is not supported when using liger kernel') diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/sampling_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/sampling_args.py new file mode 100644 index 0000000000000000000000000000000000000000..3221663d552fdb43f117179e3ca7fdfbb7b20e9f --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/sampling_args.py @@ -0,0 +1,121 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import dataclasses +import json +from dataclasses import dataclass +from datetime import datetime +from typing import List, Literal, Optional + +from swift.utils import get_logger +from .base_args import BaseArguments + +logger = get_logger() + + +@dataclass +class SamplingArguments(BaseArguments): + """A dataclass for configuring sampling parameters. + + Args: + prm_model (Optional[str]): The type of the Process Reward Model (PRM). Can be a model ID (loaded via + 'transformers' engine) or a PRM key defined in a plugin for custom inference. Defaults to None. + orm_model (Optional[str]): The type of the Outcome Reward Model (ORM). Typically a wildcard or test case, + usually defined in a plugin. Defaults to None. + sampler_type (Literal['sample', 'distill']): The type of sampling to perform. Supported types are 'sample' and + 'distill'. Defaults to 'sample'. + sampler_engine (Literal['transformers', 'lmdeploy', 'vllm', 'no', 'client']): The inference engine for the + sampling model. Supported options are 'transformers', 'lmdeploy', 'vllm', 'client', and 'no'. + Defaults to 'transformers'. + output_dir (str): The directory to save the output files. Defaults to 'sample_output'. + output_file (Optional[str]): The name of the output file. If None, a timestamp will be used as the filename. + The path should not be included, only the filename. Only the '.jsonl' format is supported. Defaults to + None. + resume (bool): Whether to resume file. Defaults to False. + override_exist_file (bool): Whether to override the output file if it already exists. This is only effective + when `output_file` is specified. Defaults to False. + num_return_sequences (int): The number of raw sequences to return from sampling. Effective for the 'sample' + `sampler_type`. Defaults to 64. + num_sampling_batch_size (int): The batch size for each sampling iteration. Defaults to 1. + num_sampling_batches (Optional[int]): The total number of batches to sample. Defaults to None. + n_best_to_keep (int): The number of best sequences to keep after evaluation. Defaults to 5. + data_range (List[int]): Specifies the data shard to process. A list of two integers `[shard_index, + num_shards]`. For example, `[1, 3]` means the dataset is split into 3 shards and this process handles the + second shard (0-indexed). Defaults to []. + temperature (float): The temperature for sampling. Defaults to 1.0. + prm_threshold (float): The threshold for the Process Reward Model (PRM). Results with a score below this + threshold will be filtered out. Defaults to 0.0. + easy_query_threshold (Optional[float]): For a single query, if the proportion of correctly sampled sequences + (as evaluated by the ORM) is greater than this threshold, the query will be discarded. This prevents overly + simple queries from appearing in the final results. Defaults to None, which disables this filter. + engine_kwargs (Optional[str]): Additional arguments to pass to the `sampler_engine`, provided as a JSON string. + For example: '{"cache_max_entry_count":0.7}'. Defaults to None. + cache_files (List[str]): A list of cache files for a two-step sampling process to avoid OOM errors. + Step 1: Set `prm_model`, and `orm_model` to None. All generated sequences are saved to a file. + Step 2: Set `sampler_engine` to 'no' and provide the output file from Step 1 to `cache_files`. + This run will perform PRM and ORM evaluation on the cached results. + Note: The `--dataset` argument must still be provided, as IDs in the cache files are MD5 hashes of the + original data and need to be linked. + """ + # rm models + prm_model: Optional[str] = None + orm_model: Optional[str] = None + + # sampler settings + sampler_type: Literal['sample', 'distill'] = 'sample' + sampler_engine: Literal['transformers', 'lmdeploy', 'vllm', 'no', 'client'] = 'transformers' + output_dir: str = 'sample_output' + output_file: Optional[str] = None + resume: bool = False + override_exist_file: bool = False + num_return_sequences: int = 64 + num_sampling_batch_size: int = 1 + num_sampling_batches: Optional[int] = None + n_best_to_keep: int = 5 + data_range: List[int] = dataclasses.field(default_factory=list) + + # generate settings + temperature: float = 1.0 + prm_threshold: float = 0.0 + easy_query_threshold: Optional[float] = None + + # engine settings + engine_kwargs: Optional[str] = None + + # Vanilla + cache_files: List[str] = dataclasses.field(default_factory=list) + + def _init_model_info(self): + if self.sampler_engine != 'client': + return super()._init_model_info() + else: + self.model_info = None + self.model_meta = None + self.task_type = 'causal_lm' + return + + def __post_init__(self): + if self.sampler_engine == 'pt': + self.sampler_engine = 'transformers' # compat swift3.x + if self.output_file is None: + now = datetime.now() + formatted_time = now.strftime('%Y-%m-%d-%H-%M-%S') + self.output_file = formatted_time + '.jsonl' + logger.info(f'Setting output_file to {self.output_file}') + else: + if '/' in self.output_file or '\\' in self.output_file: + raise ValueError(f'Please use a string prefix without directory to ' + f'`--output_file` but now is: {self.output_file}') + self.padding_side = 'left' + if self.engine_kwargs is not None: + self.engine_kwargs = json.loads(self.engine_kwargs) + else: + self.engine_kwargs = {} + + super().__post_init__() + + if self.system is not None: + self.system_message = [{ + 'role': 'system', + 'content': self.system, + }] + else: + self.system_message = [] diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/sft_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/sft_args.py new file mode 100644 index 0000000000000000000000000000000000000000..0716ec87d62085a39442456a6f02eef49d652089 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/sft_args.py @@ -0,0 +1,428 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +from dataclasses import dataclass +from transformers.utils.versions import require_version +from typing import Literal, Optional + +from swift.trainers import Seq2SeqTrainingArguments, TrainerFactory +from swift.trainers.utils import prepare_deepspeed_elastic_config +from swift.utils import (add_version_to_work_dir, get_device_count, get_logger, get_pai_tensorboard_dir, is_mp, + is_pai_training_job, is_swanlab_available, json_parse_to_dict, to_abspath) +from .base_args import BaseArguments +from .tuner_args import TunerArguments + +logger = get_logger() + + +@dataclass +class SwanlabArguments: + """Arguments for configuring Swanlab for experiment result logging. + + This dataclass stores all the configuration parameters required for initializing and using Swanlab to track + experiments. + + Args: + swanlab_token (Optional[str]): The API key for SwanLab. You can also specify it using the `SWANLAB_API_KEY` + environment variable. + swanlab_project (str): The SwanLab project, which can be created in advance on the page + [https://swanlab.cn/space/~](https://swanlab.cn/space/~) or created automatically. + The default is "ms-swift". + swanlab_workspace (Optional[str]): The SwanLab workspace. Defaults to `None`, in which case the username + associated with the API key will be used. + swanlab_exp_name (Optional[str]): The name of the experiment. If `None`, it will default to the value of the + `output_dir` argument. + swanlab_notification_method (Optional[str]): The notification method for SwanLab when training completes + or errors occur. For details, refer to [here](https://docs.swanlab.cn/plugin/notification-dingtalk.html). + Supports 'dingtalk', 'lark', 'email', 'discord', 'wxwork', 'slack'. + swanlab_webhook_url (Optional[str]): Defaults to None. The webhook URL corresponding to + SwanLab's `swanlab_notification_method`. + swanlab_secret (Optional[str]): Defaults to None. The secret corresponding to + SwanLab's `swanlab_notification_method`. + swanlab_sender_email (Optional[str]): The email address of the sender. Required when + `swanlab_notification_method` is 'email'. + swanlab_receiver_email (Optional[str]): The email address of the receiver. Required when + `swanlab_notification_method` is 'email'. + swanlab_smtp_server (Optional[str]): The SMTP server address for email notification (e.g., 'smtp.qq.com'). + swanlab_smtp_port (Optional[int]): The SMTP server port for email notification (e.g., 465). + swanlab_email_language (Optional[str]): email messages language. Supports 'zh', 'en'. The default is "zh". + swanlab_mode (Literal['cloud', 'local']): The operation mode, either 'cloud' for cloud-based logging or 'local' + for local-only logging. + """ + swanlab_token: Optional[str] = None + swanlab_project: str = 'ms-swift' + swanlab_workspace: Optional[str] = None + swanlab_exp_name: Optional[str] = None + swanlab_notification_method: Optional[str] = None + swanlab_webhook_url: Optional[str] = None + swanlab_secret: Optional[str] = None + swanlab_sender_email: Optional[str] = None + swanlab_receiver_email: Optional[str] = None + swanlab_smtp_server: Optional[str] = None + swanlab_smtp_port: Optional[int] = None + swanlab_email_language: Optional[str] = 'zh' + swanlab_mode: Literal['cloud', 'local'] = 'cloud' + + def _init_swanlab(self): + if not is_swanlab_available(): + raise ValueError('You are using swanlab as `report_to`, please install swanlab by ' + '`pip install swanlab`') + if not self.swanlab_exp_name: + self.swanlab_exp_name = self.output_dir + import swanlab + from swanlab.integration.transformers import SwanLabCallback + from transformers.integrations import INTEGRATION_TO_CALLBACK + if self.swanlab_token: + swanlab.login(self.swanlab_token) + + if self.swanlab_notification_method is not None: + from swanlab.plugin.notification import (DingTalkCallback, DiscordCallback, EmailCallback, LarkCallback, + SlackCallback, WXWorkCallback) + notification_mapping = { + 'lark': LarkCallback, + 'dingtalk': DingTalkCallback, + 'email': EmailCallback, + 'discord': DiscordCallback, + 'wxwork': WXWorkCallback, + 'slack': SlackCallback, + } + callback_cls = notification_mapping.get(self.swanlab_notification_method) + if callback_cls is None: + raise ValueError( + f'Unsupported swanlab_notification_method: "{self.swanlab_notification_method}". Supported methods' + f' are: {list(notification_mapping.keys())}') + + if self.swanlab_notification_method == 'email': + if not (self.swanlab_sender_email and self.swanlab_receiver_email and self.swanlab_smtp_server + and self.swanlab_smtp_port): + raise ValueError("When 'swanlab_notification_method' is 'email', both 'swanlab_sender_email' " + "and 'swanlab_receiver_email' and 'swanlab_smtp_server' and 'swanlab_smtp_port' " + 'must be provided.') + callback = EmailCallback( + sender_email=self.swanlab_sender_email, + receiver_email=self.swanlab_receiver_email, + password=self.swanlab_secret, + smtp_server=self.swanlab_smtp_server, + port=self.swanlab_smtp_port, + language=self.swanlab_email_language) + else: + callback = callback_cls( + webhook_url=self.swanlab_webhook_url, + secret=self.swanlab_secret, + ) + swanlab.register_callbacks([callback]) + + INTEGRATION_TO_CALLBACK['swanlab'] = SwanLabCallback( + project=self.swanlab_project, + workspace=self.swanlab_workspace, + experiment_name=self.swanlab_exp_name, + config={'UPPERFRAME': '🐦‍⬛ms-swift'}, + mode=self.swanlab_mode, + ) + + +@dataclass +class SftArguments(SwanlabArguments, TunerArguments, BaseArguments, Seq2SeqTrainingArguments): + """Arguments pertaining to the training process. + + SftArguments is a dataclass that inherits from multiple argument classes: SwanlabArguments, TunerArguments, + BaseArguments, Seq2SeqTrainingArguments. + + Args: + add_version (bool): Whether to add a versioned subdirectory like '-' to the `output_dir` to + prevent overwriting existing checkpoints. Defaults to True. + create_checkpoint_symlink (bool): Whether to create additional symbolic links for checkpoints, which can be + useful for automated training scripts. The symlinks for the best and last models will be created at + `f'{output_dir}/best'` and `f'{output_dir}/last'`, respectively. Defaults to False. + output_dir (Optional[str]): The directory to save model outputs. Defaults to 'output/'. + learning_rate (Optional[float]): The learning rate. Defaults to 1e-5 for full-parameter training and 1e-4 for + tuners like LoRA. + Note: To set a minimum learning rate (min_lr), you can pass the arguments + --lr_scheduler_type cosine_with_min_lr --lr_scheduler_kwargs '{"min_lr": 1e-6}'. + eval_strategy (Optional[str]): The evaluation strategy. By default, it aligns with `save_strategy`. It will + default to 'no' if no validation dataset is provided (i.e., `val_dataset` and `eval_dataset` are not used, + and `split_dataset_ratio` is 0). + fp16 (Optional[bool]): Defaults to None. + bf16 (Optional[bool]): Defaults to None. + max_new_tokens (int): Overrides generation parameters. The maximum number of new tokens to generate when + `predict_with_generate` is True. Defaults to 64. + temperature (float): Overrides generation parameters. The temperature for sampling when `predict_with_generate` + is True. Defaults to 0.0. + load_args (bool): Whether to load `args.json` from a saved directory when `--resume_from_checkpoint`, + `--model`, or `--adapters` is specified. For details on which keys are loaded, refer to `base_args.py`. + Defaults to `True` for inference and exporting, and `False` for training. This argument typically does not + need to be modified. + zero_hpz_partition_size (Optional[int]): A feature of ZeRO++. Enables model sharding within a node and data + sharding between nodes. If you encounter `grad_norm` NaN issues, consider trying `--torch_dtype float16`. + Defaults to None. + deepspeed_autotp_size (Optional[int]): The tensor parallelism size for DeepSpeed AutoTP. To use this, the + `--deepspeed` argument must be set to 'zero0', 'zero1', or 'zero2'. Note: This feature only supports + full-parameter fine-tuning. Defaults to None. + """ + add_version: bool = True + create_checkpoint_symlink: bool = False + + # override + output_dir: Optional[str] = None + learning_rate: Optional[float] = None + eval_strategy: Optional[str] = None # steps, epoch + fp16: Optional[bool] = None + bf16: Optional[bool] = None + + # extra + max_new_tokens: int = 64 + temperature: float = 0. + load_args: bool = False + + # zero++ + zero_hpz_partition_size: Optional[int] = None + + # auto_tp + deepspeed_autotp_size: Optional[int] = None + + # fsdp + fsdp: Optional[str] = None + + def _check_padding_free(self): + if self.padding_free or self.packing: + if self.packing: + feature = 'packing' + self.padding_free = True + else: + feature = 'padding_free' + supported_impls = ['flash_attn', 'flash_attention_2', 'flash_attention_3', 'flash_attention_4'] + if self.attn_impl not in supported_impls: + supported_impls_str = ', '.join([f'"{impl}"' for impl in supported_impls]) + raise ValueError(f'The "{feature}" feature requires a flash attention implementation. ' + f'Please use one of: {supported_impls_str}.') + + def __post_init__(self) -> None: + if self.resume_from_checkpoint: + self.resume_from_checkpoint = to_abspath(self.resume_from_checkpoint, True) + # The non-resume_only_model will have its weights loaded in the trainer. + if self.resume_only_model: + if self.tuner_type == 'full': + self.model = self.resume_from_checkpoint + else: + self.adapters = [self.resume_from_checkpoint] + BaseArguments.__post_init__(self) + self._init_override() + TunerArguments.__post_init__(self) + self._check_padding_free() + if self.vit_gradient_checkpointing is None: + self.vit_gradient_checkpointing = not self.freeze_vit + if self.optimizer is None: + if self.lorap_lr_ratio: + self.optimizer = 'lorap' + elif self.use_galore: + self.optimizer = 'galore' + + if len(self.dataset) == 0 and len(self.cached_dataset) == 0: + raise ValueError(f'self.dataset: {self.dataset}, self.cached_dataset: {self.cached_dataset}. ' + 'Please input the training dataset.') + + self._handle_pai_compat() + + self._init_deepspeed() + self._init_fsdp() + self._init_device() + + if getattr(self, 'accelerator_config', None) is None: + self.accelerator_config = {'dispatch_batches': False} + if not (self.eval_dataset or self._val_dataset_exists): + self.eval_strategy = 'no' + self.training_args = TrainerFactory.get_training_args(self) + self.training_args.remove_unused_columns = False + self._add_version() + + if 'swanlab' in self.report_to: + self._init_swanlab() + + def _init_override(self): + self._init_output_dir() + self._init_metric() + + if self.learning_rate is None: + if self.tuner_type == 'full': + self.learning_rate = 1e-5 + else: + self.learning_rate = 1e-4 + self._init_eval_strategy() + + def _init_deepspeed(self): + if self.deepspeed: + require_version('deepspeed') + if is_mp() and not self.use_ray: + raise ValueError('DeepSpeed is not compatible with `device_map`. ' + f'n_gpu: {get_device_count()}, ' + f'local_world_size: {self.local_world_size}.') + + ds_config_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'config')) + deepspeed_mapping = { + name: f'{name}.json' + for name in ['zero0', 'zero1', 'zero2', 'zero3', 'zero2_offload', 'zero3_offload'] + } + for ds_name, ds_config in deepspeed_mapping.items(): + if self.deepspeed == ds_name: + self.deepspeed = os.path.join(ds_config_folder, ds_config) + break + + self.deepspeed = json_parse_to_dict(self.deepspeed) + if self.zero_hpz_partition_size is not None: + assert 'zero_optimization' in self.deepspeed + self.deepspeed['zero_optimization']['zero_hpz_partition_size'] = self.zero_hpz_partition_size + logger.warn('If `zero_hpz_partition_size`(ZeRO++) causes grad_norm NaN, please' + ' try `--torch_dtype float16`') + if self.deepspeed_autotp_size is not None: + assert self.deepspeed is not None, ( + 'To use `deepspeed_autotp_size`, you need to additionally set the `--deepspeed` argument.') + self.deepspeed.setdefault('tensor_parallel', {})['autotp_size'] = self.deepspeed_autotp_size + self.deepspeed.setdefault('zero_optimization', {})['gather_16bit_weights_on_model_save'] = True + if 'deepspeed_elastic' in set(getattr(self, 'callbacks', []) or []): + prepare_deepspeed_elastic_config(self) + logger.info(f'Using deepspeed: {self.deepspeed}') + + def _init_fsdp(self): + if not self.fsdp: + self.fsdp = [] + return + + if is_mp() and not self.use_ray: + raise ValueError('FSDP2 is not compatible with `device_map`. ' + f'n_gpu: {get_device_count()}, ' + f'local_world_size: {self.local_world_size}.') + if self.deepspeed: + raise ValueError('FSDP2 is not compatible with DeepSpeed.') + + fsdp_config_folder = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'config')) + + # FSDP2 preset configurations + fsdp_mapping = { + 'fsdp2': 'fsdp2.json', + } + + fsdp_config_path = self.fsdp + for fsdp_name, fsdp_config in fsdp_mapping.items(): + if self.fsdp == fsdp_name: + fsdp_config_path = os.path.join(fsdp_config_folder, fsdp_config) + break + + fsdp_config_dict = json_parse_to_dict(fsdp_config_path) + + # Extract fsdp string options (e.g., "full_shard auto_wrap offload") + fsdp_options = fsdp_config_dict.get('fsdp', 'full_shard auto_wrap') + self.fsdp = fsdp_options + + # Extract fsdp_config dict + self.fsdp_config = fsdp_config_dict.get('fsdp_config', {}) + + # Set FSDP_VERSION environment variable for accelerate to recognize FSDP2 + fsdp_version = self.fsdp_config.get('fsdp_version', 2) + os.environ['FSDP_VERSION'] = str(fsdp_version) + + # Set environment variable to optimize NCCL memory usage + if 'TORCH_NCCL_AVOID_RECORD_STREAMS' not in os.environ: + os.environ['TORCH_NCCL_AVOID_RECORD_STREAMS'] = '1' + + # Check FSDP2 compatibility with other training arguments + self._check_fsdp2_compatibility() + + logger.info(f'Using FSDP2: fsdp={self.fsdp}, fsdp_config={self.fsdp_config}') + + def _check_fsdp2_compatibility(self): + """Check for incompatible argument combinations with FSDP2. + + FSDP2 has several known limitations: + 1. save_only_model=True + SHARDED_STATE_DICT: Can't save only model weights with sharded state dict + 2. gradient_checkpointing=True: Should use activation_checkpointing in fsdp_config instead + """ + state_dict_type = self.fsdp_config.get('state_dict_type', 'SHARDED_STATE_DICT') + + # Check 1: save_only_model + SHARDED_STATE_DICT + if getattr(self, 'save_only_model', False) and 'SHARDED' in state_dict_type.upper(): + raise ValueError( + 'FSDP2 with SHARDED_STATE_DICT is not compatible with save_only_model=True. ' + 'Either set save_only_model=False, or change state_dict_type to FULL_STATE_DICT in fsdp_config. ' + 'Note: FULL_STATE_DICT requires more memory and is slower.') + + # Check 2: gradient_checkpointing should be disabled, use activation_checkpointing instead + if getattr(self, 'gradient_checkpointing', False): + activation_checkpointing = self.fsdp_config.get('activation_checkpointing', False) + if activation_checkpointing: + logger.warning('Both gradient_checkpointing and fsdp_config.activation_checkpointing are enabled. ' + 'For FSDP2, it is recommended to use only activation_checkpointing in fsdp_config. ' + 'Disabling gradient_checkpointing automatically.') + self.gradient_checkpointing = False + else: + logger.warning( + 'gradient_checkpointing is enabled with FSDP2. ' + 'For better performance, consider using activation_checkpointing in fsdp_config instead. ' + 'Add "activation_checkpointing": true to your fsdp_config.') + + def _handle_pai_compat(self) -> None: + if not is_pai_training_job(): + return + + logger.info('Handle pai compat...') + pai_tensorboard_dir = get_pai_tensorboard_dir() + if self.logging_dir is None and pai_tensorboard_dir is not None: + self.logging_dir = pai_tensorboard_dir + logger.info(f'Setting args.logging_dir: {self.logging_dir}') + self.add_version = False + logger.info(f'Setting args.add_version: {self.add_version}') + + def _add_version(self): + """Prepare the output_dir""" + if self.add_version: + self.output_dir = add_version_to_work_dir(self.output_dir) + logger.info(f'output_dir: {self.output_dir}') + + if self.logging_dir is None: + self.logging_dir = f'{self.output_dir}/runs' + + self.logging_dir = to_abspath(self.logging_dir) + # transformers>=5.15 dropped `TrainingArguments.logging_dir`; its TensorBoardCallback reads the log + # dir from this environment variable instead. Harmless on older versions, which ignore it. + os.environ.setdefault('TENSORBOARD_LOGGING_DIR', self.logging_dir) + os.makedirs(self.output_dir, exist_ok=True) + + if self.run_name is None: + self.run_name = self.output_dir + + self.training_args.output_dir = self.output_dir + self.training_args.run_name = self.run_name + self.training_args.logging_dir = self.logging_dir + + def _init_output_dir(self): + if self.output_dir is None: + self.output_dir = f'output/{self.model_suffix}' + self.output_dir = to_abspath(self.output_dir) + + def _init_eval_strategy(self): + if self.eval_strategy is None: + self.eval_strategy = self.save_strategy + if self.eval_strategy == 'no': + self.eval_steps = None + if self.split_dataset_ratio > 0: + self.split_dataset_ratio = 0. + logger.info(f'Setting args.split_dataset_ratio: {self.split_dataset_ratio}') + elif self.eval_strategy == 'steps' and self.eval_steps is None: + self.eval_steps = self.save_steps + self.evaluation_strategy = self.eval_strategy + + def _init_metric(self): + if self.eval_metric is None: + if self.task_type == 'causal_lm' and self.predict_with_generate: + self.eval_metric = 'nlg' + elif self.task_type == 'embedding': + self.eval_metric = 'infonce' if self.loss_type == 'infonce' else 'paired' + elif self.task_type in {'reranker', 'generative_reranker'}: + self.eval_metric = 'reranker' + if self.eval_metric == 'nlg': + require_version('jieba', 'Setting `--eval_metric nlg` requires installing the jieba dependency.') + self._init_metric_for_best_model() + + def _init_metric_for_best_model(self): + if self.metric_for_best_model is None: + self.metric_for_best_model = 'rouge-l' if self.predict_with_generate else 'loss' + if self.greater_is_better is None and self.metric_for_best_model is not None: + self.greater_is_better = 'loss' not in self.metric_for_best_model diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/tuner_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/tuner_args.py new file mode 100644 index 0000000000000000000000000000000000000000..b742cec56a287cb5c56261296f432de8506d90c1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/tuner_args.py @@ -0,0 +1,220 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import dataclass, field +from transformers.utils import strtobool +from typing import List, Literal, Optional + +from swift.utils import get_logger + +logger = get_logger() + + +@dataclass +class TunerArguments: + """ + TunerArguments is a dataclass that holds configuration for various tuners. + + Args: + freeze_parameters (List[str]): A list of prefixes for parameters that should be frozen during training. + Defaults to an empty list `[]`. + freeze_parameters_regex (Optional[str]): A regular expression to match the names of parameters that should be + frozen. Defaults to `None`. + freeze_parameters_ratio (float): The ratio of parameters to freeze, starting from the bottom layers upwards + (from 0.0 to 1.0). Setting this to 1.0 freezes all model parameters, which can be useful when selectively + unfreezing specific parameters with `trainable_parameters`. Defaults to 0.0. + trainable_parameters (List[str]): A list of prefixes for parameters that should be made explicitly trainable. + Defaults to an empty list `[]`. + trainable_parameters_regex (Optional[str]): A regular expression to match the names of parameters that should + be made explicitly trainable. Defaults to `None`. + Note on parameter freezing priority: The `trainable_*` arguments have higher priority than the `freeze_*` + arguments. The freezing logic is applied as follows: + Firstly, all parameters are set to trainable. + Then, `freeze_parameters`, `freeze_parameters_regex`, and `freeze_parameters_ratio` are applied to freeze + parts of the model. + Finally, `trainable_parameters` and `trainable_parameters_regex` are used to unfreeze specific parameters, + ensuring they are trainable regardless of the freezing rules. + + freeze_llm (bool): For multi-modal models only. If `True`, it affects the Large Language Model (LLM) part. + In full fine-tuning, this freezes the LLM weights. In LoRA training with `target_modules=['all-linear']`, + this prevents adding LoRA modules to the LLM. Defaults to `False`. + freeze_vit (bool): For multi-modal models only. If `True`, it affects the Vision/Audio Transformer (ViT) part. + In full fine-tuning, this freezes the ViT weights. In LoRA training with `target_modules=['all-linear']`, + this prevents adding LoRA modules to the ViT. Note: 'vit' can refer to `vision_tower` and `audio_tower`. + Defaults to `True`. + freeze_aligner (bool): For multi-modal models only. If `True`, it affects the aligner (projector) part. + In full fine-tuning, this freezes the aligner weights. In LoRA training with + `target_modules=['all-linear']`, this prevents adding LoRA modules to the aligner. Defaults to `True`. + + target_modules (List[str]): List of target modules for tuning. Default is ['all-linear']. + target_regex (Optional[str]): Regular expression to match target modules. Default is None. + target_parameters (Optional[List[str]]): A list of parameter names to be replaced by LoRA modules. This is + similar to `target_modules` but targets parameters directly, which is useful for layers like MoE that use + `nn.Parameter` instead of `nn.Linear`. Requires `peft>=0.17.0`. Defaults to `None`. + modules_to_save (List[str]): List of modules to save. Default is an empty list. + + lora_rank (int): Rank for LoRA. Default is 8. + lora_alpha (int): Alpha value for LoRA. Default is 32. + lora_dropout (float): Dropout rate for LoRA. Default is 0.05. + lora_bias (Literal['none', 'all']): The possible values are 'none' and 'all'. If set to 'all', all biases + will be trainable. Default is 'none'. + lora_dtype (Literal): Data type for LoRA. Default is 'AUTO'. Allowed values are 'fp16', 'bf16', 'fp32', 'AUTO'. + lorap_lr_ratio (float): Learning rate ratio for LoRA. Default is None. + use_rslora (bool): Flag to indicate if RSLora is used. Default is False. + use_dora (bool): Flag to indicate if Dora is used. Default is False. + + lora_ga_batch_size (int): Batch size used for estimating gradients during initialization in LoRA-GA. Default + value is 2. + lora_ga_iters (int): Number of iterations for estimating gradients during initialization in LoRA-GA. Default + value is 2. + lora_ga_max_length (int): Maximum input length for estimating gradients during initialization in LoRA-GA. + Default value is 1024. + lora_ga_direction (str): Initial direction used for gradient estimation during initialization in LoRA-GA. + Default value is `ArB2r`. Allowed: `ArBr`, `A2rBr`, `ArB2r`, and `random`. + lora_ga_scale (str): The scaling method for initialization in LoRA-GA. + Default value is `stable`. Allowed values are: `gd`, `unit`, `stable`, and `weightS`. + lora_ga_stable_gamma (int): The gamma value when choosing `stable` scaling for initialization. Default + value is 16. + + init_weights (str): The method for initializing adapter weights. For LoRA, options include 'true', 'false', + 'gaussian', 'pissa', 'pissa_niter_[number of iters]', 'olora', 'loftq', and 'lora-ga'. For BoNE, + options are 'true', 'false', and 'bat'. Defaults to 'true'. + + fourier_n_frequency (int): Number of frequencies for FourierFT. Default is 2000. + fourier_scaling (float): Scaling factor for FourierFT. Default is 300.0. + + boft_block_size (int): Block size for BOFT. Default is 4. + boft_block_num (int): Number of blocks for BOFT. Default is 0. + boft_n_butterfly_factor (int): Butterfly factor for BOFT. Default is 1. + boft_dropout (float): Dropout rate for BOFT. Default is 0.0. + + vera_rank (int): Rank for Vera. Default is 256. + vera_projection_prng_key (int): PRNG key for Vera projection. Default is 0. + vera_dropout (float): Dropout rate for Vera. Default is 0.0. + vera_d_initial (float): Initial value for Vera D. Default is 0.1. + + adapter_act (str): Activation function for adapter. Default is 'gelu'. + adapter_length (int): Length of the adapter. Default is 128. + + adalora_target_r (int): Target rank for AdaLoRA. Default is 8. + adalora_init_r (int): Initial rank for AdaLoRA. Default is 12. + adalora_tinit (int): Initial T value for AdaLoRA. Default is 100. + adalora_tfinal (int): Final T value for AdaLoRA. Default is 1000. + adalora_deltaT (int): Delta T value for AdaLoRA. Default is 10. + adalora_beta1 (float): Beta1 value for AdaLoRA. Default is 0.85. + adalora_beta2 (float): Beta2 value for AdaLoRA. Default is 0.85. + adalora_orth_reg_weight (float): Orthogonal regularization weight for AdaLoRA. Default is 0.5. + + llamapro_num_new_blocks (int): Number of new blocks for LLaMAPro. Default is 4. + llamapro_num_groups (Optional[int]): Number of groups for LLaMAPro. Default is None. + + reft_layer_key (Optional[str]): Key identifier for ReFT layer. Default is None. + reft_layers (Optional[List[int]]): List of layers involved in ReFT. Default is None. + reft_rank (int): Rank parameter for ReFT. Default is 4. + reft_intervention_type (Literal): Type of intervention for ReFT. Default is 'LoreftIntervention'. + reft_args (Optional[str]): Additional arguments for ReFT. Default is None. + """ + # full + freeze_parameters: List[str] = field(default_factory=list) + freeze_parameters_regex: Optional[str] = None + freeze_parameters_ratio: float = 0. # 0 ~ 1 + trainable_parameters: List[str] = field(default_factory=list) + trainable_parameters_regex: Optional[str] = None + # lora or full + freeze_llm: bool = False + freeze_vit: bool = True + freeze_aligner: bool = True + # tuners + target_modules: List[str] = field(default_factory=lambda: ['all-linear']) + target_regex: Optional[str] = None + target_parameters: Optional[List[str]] = None + # e.g. ['wte', 'ln_1', 'ln_2', 'ln_f', 'lm_head'] + modules_to_save: List[str] = field(default_factory=list) + + # lora + lora_rank: int = 8 + lora_alpha: int = 32 + lora_dropout: float = 0.05 + lora_bias: Literal['none', 'all'] = 'none' + lora_dtype: Literal['float16', 'bfloat16', 'float32', None] = None + lorap_lr_ratio: Optional[float] = None + use_rslora: bool = False + use_dora: bool = False + + # lora_ga + lora_ga_batch_size: int = 2 + lora_ga_iters: int = 2 + lora_ga_max_length: int = 1024 + lora_ga_direction: str = 'ArB2r' + lora_ga_scale: str = 'stable' + lora_ga_stable_gamma: int = 16 + + # Lora: Literal['gaussian', 'pissa', 'pissa_niter_[number of iters]', 'olora', 'loftq', 'true', 'false', 'lora-ga'] + # Bone: Literal['bat', 'true', 'false'] + init_weights: str = 'true' + + # fourierft + fourier_n_frequency: int = 2000 + fourier_scaling: float = 300.0 + + # BOFT + boft_block_size: int = 4 + boft_block_num: int = 0 + boft_n_butterfly_factor: int = 1 + boft_dropout: float = 0.0 + + # Vera + vera_rank: int = 256 + vera_projection_prng_key: int = 0 + vera_dropout: float = 0.0 + vera_d_initial: float = 0.1 + + # adapter + adapter_act: str = 'gelu' + adapter_length: int = 128 + + # adalora + adalora_target_r: int = 8 + adalora_init_r: int = 12 + adalora_tinit: int = 0 + adalora_tfinal: int = 0 + adalora_deltaT: int = 1 + adalora_beta1: float = 0.85 + adalora_beta2: float = 0.85 + adalora_orth_reg_weight: float = 0.5 + + # llamapro + llamapro_num_new_blocks: int = 4 + llamapro_num_groups: Optional[int] = None + + # reft + reft_layer_key: Optional[str] = None + reft_layers: Optional[List[int]] = None + reft_rank: int = 4 + reft_intervention_type: Literal['NoreftIntervention', 'LoreftIntervention', 'ConsreftIntervention', + 'LobireftIntervention', 'DireftIntervention', + 'NodireftIntervention'] = 'LoreftIntervention' + reft_args: Optional[str] = None + + def __post_init__(self): + if isinstance(self.init_weights, str) and self.init_weights.lower() in {'true', 'false'}: + self.init_weights = bool(strtobool(self.init_weights)) + self._init_multimodal_full() + if self.target_regex: + self.target_modules = self.target_regex + + def _init_multimodal_full(self): + model_arch = self.model_meta.model_arch + if not self.model_meta.is_multimodal or not model_arch or self.tuner_type != 'full': + return + if self.freeze_llm: + self.freeze_parameters += model_arch.language_model + if self.freeze_vit: + self.freeze_parameters += model_arch.vision_tower + if self.freeze_aligner: + self.freeze_parameters += model_arch.aligner + else: + self.trainable_parameters += model_arch.aligner + self.freeze_parameters += model_arch.generator + if self.freeze_parameters: + logger.info(f'freeze_parameters: {self.freeze_parameters}') + if self.trainable_parameters: + logger.info(f'additional trainable_parameters: {self.trainable_parameters}') diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/arguments/webui_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/webui_args.py new file mode 100644 index 0000000000000000000000000000000000000000..105446c608412b5bd3a710cba3709b441d9b6b03 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/arguments/webui_args.py @@ -0,0 +1,18 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import dataclass + + +@dataclass +class WebUIArguments: + """A dataclass for web UI configuration arguments. + + Args: + server_name (str): The hostname or IP address to be bound to the Web UI server. Defaults to '0.0.0.0'. + server_port (int): The port number to be bound to the Web UI server. Defaults to 7860. + share (bool): Whether to create a public, shareable link for the web UI. Defaults to False. + lang (str): The language for the web UI, chosen from {'zh', 'en'}. Defaults to 'zh'. + """ + server_name: str = '0.0.0.0' + server_port: int = 7860 + share: bool = False + lang: str = 'zh' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..88650b6bcfcded3f553d286afeda59cbb8b31238 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/__init__.py @@ -0,0 +1,2 @@ +from .base import TrainerCallback +from .mapping import callbacks_map diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/activation_cpu_offload.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/activation_cpu_offload.py new file mode 100644 index 0000000000000000000000000000000000000000..7ff467c20dc78fcd9f53219b5f07cd48dc45f951 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/activation_cpu_offload.py @@ -0,0 +1,610 @@ +"""Functionality for CPU offloading of tensors saved for backward pass.""" +import functools +import torch +from torch.distributed.fsdp import FSDPModule as FSDP2 +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from transformers.trainer_callback import TrainerControl, TrainerState +from transformers.training_args import TrainingArguments +from transformers.utils import is_torch_npu_available +from typing import Any, Optional + +from swift.utils import get_logger +from .base import TrainerCallback + +logger = get_logger() + +is_cuda_available = torch.cuda.is_available() +is_npu_available = is_torch_npu_available() + + +def _get_unique_tensor_key(tensor): + key = (tensor.untyped_storage().data_ptr() + tensor.storage_offset(), tensor.dtype) + return key + + +def get_device_name() -> str: + """Function that gets the torch.device based on the current machine. + This currently only supports CPU, CUDA, NPU. + Returns: + device + """ + if is_cuda_available: + device = 'cuda' + elif is_npu_available: + device = 'npu' + else: + device = 'cpu' + return device + + +class FSDPParameterFilter: + + def __init__(self): + self.model_parameters_storage = set() + + def __call__(self, tensor): + return tensor.untyped_storage().data_ptr() not in self.model_parameters_storage + + def update_model_parameters(self, model): + new_storage = set() + for p in model.parameters(): + new_storage.add(p.data.untyped_storage().data_ptr()) + self.model_parameters_storage = new_storage + + +def get_torch_device() -> Any: + """Return the corresponding torch attribute based on the device type string. + Returns: + module: The corresponding torch device namespace, or torch.cuda if not found. + """ + device_name = get_device_name() + try: + return getattr(torch, device_name) + except AttributeError: + logger.warning(f"Device namespace '{device_name}' not found in torch, try to load torch.cuda.") + return torch.cuda + + +class CpuOffloadHookWithOffloadHandler: + """Context-manager that offloads/recovers tensors through an offload hander. + + The hook just offloads/recovers the tensor object to the handler through `tensor_push` + and `tensor_pop` interface. How the offload-handler manages the offloading, recovering + or prefetching timing is transparent to this hook. + """ + + def __init__( + self, + offload_handler: 'OffloadHandler', + handler_extra_kwargs: Optional[dict[str, Any]] = None, + ) -> None: + if handler_extra_kwargs is None: + handler_extra_kwargs = {} + self.offload_handler: OffloadHandler = offload_handler + self.handler_extra_kwargs: dict[str, Any] = handler_extra_kwargs + self.inside_context = False + + def __enter__(self): + self.inside_context = True + torch._C._autograd._push_saved_tensors_default_hooks(self.on_save_for_backward, self.on_get_saved_tensor) + + def __exit__(self, *args: Any): + self.inside_context = False + torch._C._autograd._pop_saved_tensors_default_hooks() + + def on_save_for_backward(self, tensor: torch.Tensor) -> Any: + retrieve_identifier = self.offload_handler.tensor_push(tensor, **self.handler_extra_kwargs) + return retrieve_identifier + + def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: + tensor = self.offload_handler.tensor_pop(saved_state, **self.handler_extra_kwargs) + return tensor + + +class OffloadHandler: + """A base class for CPU offload-handler.""" + + def __init__(self) -> None: + pass + + def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: + """Tensor push.""" + raise NotImplementedError( + '`tensor_push is not implented in OffloadHandler class. Inherit this class and implement your ' + 'custom tensor_push.') + + def tensor_pop(self, tensor_tag: Any, **kwargs): + """Tensor pop.""" + raise NotImplementedError( + '`tensor_pop is not implented in OffloadHandler class. Inherit this class and implement your ' + 'custom tensor_pop.') + + +class GroupCommitFunction(torch.autograd.Function): + """this is a dummy op with output identical to input. + However, it is necessary for marking a timepoint for offload handler to + accomplish all synchronizations. Implementing it as a function is necessary + because we need to actions in both forward and backward. + """ + + @staticmethod + def forward(ctx, tensor, cpu_offload_handler): + # pylint: disable=missing-function-docstring + cpu_offload_handler.on_group_commit_forward() + ctx.cpu_offload_handler = cpu_offload_handler + # return the identical tensor + return tensor + + @staticmethod + def backward(ctx, grad_output): + # pylint: disable=missing-function-docstring + cpu_offload_handler = ctx.cpu_offload_handler + cpu_offload_handler.on_group_commit_backward() + return grad_output, None + + +group_prefetch_offload_commit = GroupCommitFunction.apply + + +class SynchronizedGroupOffloadHandler(OffloadHandler): + """Offload Handler that offloads/reloads in a synchronized way. + The device-to-host and host-to-device copying happen in the same stream + as the computation kernels, thus the copying will block computation. + """ + + def __init__(self, num_offload_group, tensor_need_offloading_checker=(lambda _: True)) -> None: + super().__init__() + + self.num_offload_group = num_offload_group + self.tensor_need_offloading_checker = tensor_need_offloading_checker + + self.groupid_reset() + + def groupid_reset(self): + """Groupid reset.""" + # Data structures to label saved tensors and book-keep their cpu copies. + # Currently, on push, create a new cpu tensor and copies; on pop, copies + # the tensor back to gpu and deletes the cpu tensor. + # These will increment whenever `group_commit()` is invoked + self.current_group, self.tensor_count_current_group = (0, 0) + self.torch_tensor_count = 0 + self.tensor_tag_to_state = {} + + def on_group_commit_forward(self): + """On group commit forward.""" + # finishing up with updating current group and tensor count + self.current_group += 1 # increment + self.tensor_count_current_group = 0 # reset + + def on_group_commit_backward(self): + """On group commit backward.""" + self.current_group -= 1 + assert self.current_group >= 0 + + @staticmethod + def offload(src_tensor, pin_memory=True): + """Offload.""" + # NPU doesn't fully support async H2D/D2H with pinned memory; use sync copy. + if is_npu_available: + cpu_backup = torch.empty( + src_tensor.size(), + dtype=src_tensor.dtype, + layout=src_tensor.layout, + device='cpu', + pin_memory=False, + ) + cpu_backup.copy_(src_tensor, non_blocking=False) + else: + cpu_backup = torch.empty( + src_tensor.size(), + dtype=src_tensor.dtype, + layout=src_tensor.layout, + device='cpu', + pin_memory=pin_memory, + ) + cpu_backup.copy_(src_tensor, non_blocking=True) + state = (src_tensor.device, cpu_backup) + return state + + @staticmethod + def reload(state, non_blocking=None): + """Reload.""" + dev, cpu_backup = state + if non_blocking is None: + non_blocking = cpu_backup.is_pinned() + return cpu_backup.to(dev, non_blocking=non_blocking) + + def tensor_push(self, tensor: torch.Tensor, **kwargs): + """Tensor push.""" + # obtain a unique tensor tag + tensor_tag = (self.current_group, self.tensor_count_current_group) + self.tensor_count_current_group += 1 + assert tensor_tag not in self.tensor_tag_to_state + if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(tensor): + state = SynchronizedGroupOffloadHandler.offload(tensor) + self.tensor_tag_to_state[tensor_tag] = state + else: + # will be offloaded together after group commit + self.tensor_tag_to_state[tensor_tag] = tensor + + return tensor_tag + + def tensor_pop(self, tensor_tag, **kwargs): + """Tensor pop.""" + assert tensor_tag in self.tensor_tag_to_state + state = self.tensor_tag_to_state.pop(tensor_tag) + if isinstance(state, tuple): + tensor = SynchronizedGroupOffloadHandler.reload(state) + else: + tensor = state + return tensor + + +class AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandler): + """Compared to synchronize, this uses more memory because of the buffer but + achieves better performance due to the overlapping. D2h and h2d copying are + completely hidden behind computation if computation time of a layer is longer + than host-device communication time. Bulk offloading with delay and bulk reloading + with prefetch are implemented.""" + + def __init__( + self, + num_offload_group, # must be <= actual number of groups (number of commits) + num_model_group, + tensor_need_offloading_checker=(lambda t: True), + ) -> None: + super().__init__( + num_offload_group=num_offload_group, + tensor_need_offloading_checker=tensor_need_offloading_checker, + ) + # Number of layers in the model + self.num_layers = num_model_group + # Data Structure to maintain reference to activation tensors + self.tensor_tag_to_buf = {} + # Tracking the number of layers offloaded + self.offloaded_group_count = 0 + # Core data structure that decides the window for offloading + self.layer_window_map = {} + self.group_offload_mapping = {} + + # Logic to make offloading load balance across computation + # for optimal CPU/GPU interconnect usage + constant = 0 + for i in range(self.num_offload_group): + self.layer_window_map[i] = ((self.num_layers // self.num_offload_group) * (i + 1)) - 1 + if i < (self.num_layers % self.num_offload_group): + self.layer_window_map[i] += i + 1 + constant = i + 1 + else: + self.layer_window_map[i] += constant + + # allocate streams and events for synchronization + self.d2h_stream = get_torch_device().Stream() + self.h2d_stream = get_torch_device().Stream() + + def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: + torch_stray_tensor = isinstance( + tensor, + torch._subclasses.fake_tensor.FakeTensor | torch._subclasses.functional_tensor.FunctionalTensor, + ) + need_offload = not torch_stray_tensor + need_offload = need_offload and self.tensor_need_offloading_checker(tensor) + + if need_offload: + # obtain a unique tensor tag + tensor_tag = (self.current_group, self.tensor_count_current_group) + self.tensor_count_current_group += 1 + + assert tensor_tag not in self.tensor_tag_to_state + self.tensor_tag_to_state[tensor_tag] = tensor + + if self.current_group < self.num_offload_group: + self.tensor_tag_to_buf[tensor_tag] = tensor + else: + tensor_tag = tensor + return tensor_tag + + def tensor_pop(self, tensor_tag, **kwargs): + """Tensor pop.""" + if isinstance(tensor_tag, torch.Tensor): + return tensor_tag + assert tensor_tag in self.tensor_tag_to_state + tensor = self.tensor_tag_to_state.pop(tensor_tag) + self.tensor_tag_to_buf.pop(tensor_tag, None) + + # the tensor should have been copied back in on_group_commit_backward() + # which invokes bulk_reload_group. + assert not isinstance(tensor, tuple) + return tensor + + def bulk_offload_group(self, group_to_offload): + """Bulk offload group.""" + offload_mapping = {} + offload_size = 0 + with get_torch_device().stream(self.d2h_stream): + for tensor_tag, state in self.tensor_tag_to_state.items(): + group_id, _ = tensor_tag + if group_id == group_to_offload: + assert not isinstance(state, tuple) + key = _get_unique_tensor_key(state) + if key not in offload_mapping: + offload_mapping[key] = state + # if offload, return the reference to cpu copy + self.tensor_tag_to_state[tensor_tag] = (key, state.shape) + for key, tensor in offload_mapping.items(): + state = SynchronizedGroupOffloadHandler.offload(tensor) + offload_size += tensor.numel() * tensor.element_size() + offload_mapping[key] = state + + self.group_offload_mapping[group_to_offload] = offload_mapping + + def synchronize_on_group_commit_forward(self, current_group): + """Synchronize on group commit forward.""" + + # For the first group, kickstart the offload after we have + # the first compute completion + if current_group == 0: + self.d2h_stream.wait_stream(get_torch_device().current_stream()) + self.bulk_offload_group(current_group) + + # Window map data structure helps us synchronize based on number + # of layers offloaded + if self.layer_window_map[self.offloaded_group_count] == current_group: + # Stream synchronization both ways + self.d2h_stream.wait_stream(get_torch_device().current_stream()) + get_torch_device().current_stream().wait_stream(self.d2h_stream) + + # Time to free the activation memory after usage + for tensor_tag, _ in self.tensor_tag_to_buf.items(): + if tensor_tag[0] == self.offloaded_group_count: + self.tensor_tag_to_buf[tensor_tag] = None + + # Time to offload the next group + if self.offloaded_group_count < (self.num_offload_group - 1): + self.bulk_offload_group(self.offloaded_group_count + 1) + + # Increment the offload group count to keep track + self.offloaded_group_count += 1 + + def on_group_commit_forward(self): + """This function will cause host device synchronization""" + # handle synchronization events + self.synchronize_on_group_commit_forward(self.current_group) + + super().on_group_commit_forward() + + @torch.no_grad + def bulk_reload_group(self, group_to_reload): + """Bulk reload group.""" + assert group_to_reload < self.num_offload_group + + with get_torch_device().stream(self.h2d_stream): + # move back tensors + offload_mapping = self.group_offload_mapping.pop(group_to_reload) + assert offload_mapping is not None + for key, state in offload_mapping.items(): + offload_mapping[key] = SynchronizedGroupOffloadHandler.reload(state) + for tensor_label, state in self.tensor_tag_to_state.items(): + group_id, _ = tensor_label + if group_id == group_to_reload and not isinstance(state, torch.Tensor): + assert isinstance(state, tuple), f'{group_id} {state}' + key, shape = state + recovered_tensor = offload_mapping[key].view(shape) + self.tensor_tag_to_state[tensor_label] = recovered_tensor + + def on_group_commit_backward(self): + # first decrement the current group. + # after last commit in forward, the group will +1; in backward it -1. + # Finally it should be decremented to 0. + self.current_group -= 1 + assert self.current_group >= 0 + + # Layer window data structure helps us to reload at right times + if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group: + # Stream synchronization both ways + self.h2d_stream.wait_stream(get_torch_device().current_stream()) + get_torch_device().current_stream().wait_stream(self.h2d_stream) + + # Time to reload the next group + self.bulk_reload_group(self.offloaded_group_count - 1) + + # Decrease the offloading group counter + self.offloaded_group_count -= 1 if self.offloaded_group_count > 1 else 0 + + # Last group computation needs to wait till all the reloads complete + if self.current_group == 0: + get_torch_device().current_stream().wait_stream(self.h2d_stream) + self.offloaded_group_count = 0 + + +def get_activation_offload_context(num_layers: int = 1, + model_layers: int = 1, + tensor_need_offloading_checker=(lambda t: True)): + cpu_offload_handler = AsyncDoubleBufferGroupOffloadHandler( + num_offload_group=num_layers, + num_model_group=model_layers, + tensor_need_offloading_checker=tensor_need_offloading_checker, + ) + + def group_prefetch_offload_commit_async(tensor): + return group_prefetch_offload_commit(tensor, cpu_offload_handler) + + return ( + CpuOffloadHookWithOffloadHandler(offload_handler=cpu_offload_handler), + group_prefetch_offload_commit_async, + ) + + +class ActivationHandler: + + def __init__(self, offload_ctx, sync_func, tensor_filter, enable_ckpt): + self._offload_ctx = offload_ctx + self._sync_func = sync_func + self._enable_ckpt = enable_ckpt + self._tensor_filter = tensor_filter + if enable_ckpt: + self.checkpoint_fn = functools.partial( + torch.utils.checkpoint.checkpoint, + use_reentrant=True, + ) + + def pre_forward(self, module): + if module.training: + self._offload_ctx.__enter__() + self._tensor_filter.update_model_parameters(module) + + def post_forward(self, module): + if module.training: + self._offload_ctx.__exit__(None, None, None) + + def _pack_kwargs(self, *args, **kwargs): + kwarg_keys = [] + flat_args = list(args) + for k, v in kwargs.items(): + kwarg_keys.append(k) + flat_args.append(v) + + return tuple(flat_args), tuple(kwarg_keys) + + def _unpack_kwargs(self, flat_args, kwarg_keys): + assert len(kwarg_keys) <= len(flat_args), f'too many keys {len(kwarg_keys)} vs. {len(flat_args)}' + if len(kwarg_keys) == 0: + return flat_args, {} + args = flat_args[:-len(kwarg_keys)] + kwargs = dict(zip(kwarg_keys, flat_args[-len(kwarg_keys):], strict=True)) + return args, kwargs + + def _ckpt_forward(self, forward_method, *args, **kwargs): + flat_args, kwarg_keys = self._pack_kwargs(*args, **kwargs) + + def my_function(*inputs): + # unpack back into args and kwargs + unpacked_args, unpacked_kwargs = self._unpack_kwargs(inputs, kwarg_keys) + # run original module + return forward_method(*unpacked_args, **unpacked_kwargs) + + return self.checkpoint_fn( + my_function, + *flat_args, + ) + + def forward(self, module, forward_method, *args, **kwargs): + if not module.training: + return forward_method(*args, **kwargs) + if not self._enable_ckpt: + ret = forward_method(*args, **kwargs) + else: + ret = self._ckpt_forward(forward_method, *args, **kwargs) + binded_tensor = ret + if isinstance(ret, tuple): + binded_tensor = ret[0] + binded_tensor = self._sync_func(binded_tensor) + final_ret = binded_tensor + if isinstance(ret, tuple): + final_ret = (final_ret, ) + ret[1:] + return final_ret + + def wrap_module_forward_method(self, module): + orig_method = module.forward + handler = self + + @functools.wraps(orig_method) + def wrapped_method(model_self, *args, **kwargs): + if not model_self.training or not torch.is_grad_enabled(): + return orig_method(*args, **kwargs) + handler.pre_forward(model_self) + try: + return handler.forward(model_self, orig_method, *args, **kwargs) + finally: + handler.post_forward(model_self) + + module.forward = wrapped_method.__get__(module, type(module)) + + +def enable_activation_offloading(model, strategy, enable_ckpt=False): + """ + Enable activation offloading for the model. It groups activations by TransformerLayer and offloads activation + groups asynchronously. This means that the offloading of the i-th activation group and the computation of the i+1-th + activation group happen at the same time, and there are at most two activation groups in GPU memory. + + Args: + model: the model to enable activation offloading + strategy: the training strategy of the model, such as "fsdp" + enable_ckpt: whether activation checkpointing(also called gradient checkpointing) has been enabled for the model + + Note: + For best efficiency, activation offloading is usually combined with activation checkpointing. However, this + implementation of activation offloading is conflicted with the implementation of activation checkpointing in + some training strategies. This function resolves this conflict, and therefore requires the "strategy" and + "enable_ckpt" arguments. + + Returns: + + """ + + assert strategy == 'fsdp' or strategy == 'fsdp2', 'activation offloading only supports fsdp strategy' + layers = [] + + def get_layers(module): + for name, child in module.named_children(): + if not isinstance(child, FSDP | FSDP2): + get_layers(child) + else: + wrapped_module = child + if isinstance(child, FSDP): + wrapped_module = child._fsdp_wrapped_module + # In some cases, torch.nn.Embedding is wrapped with FSDP alone. However, the activation + # size of torch.nn.Embedding is small, so it's not necessary to offload it. + if not isinstance(wrapped_module, torch.nn.Embedding): + layers.append(child) + + get_layers(model) + if len(layers) < 3: + logger.warning(f'Find only {len(layers)} fsdp layers, not necessary to enable async activation offloading') + return + + tensor_filter = FSDPParameterFilter() + context, sync_func = get_activation_offload_context(len(layers) - 1, len(layers), tensor_filter) + if enable_ckpt: + # The implementation of activation checkpointing in transformers library is incompatible with + # activation offloading, + # so it will be disabled, but this implementation supports another version of activation checkpointing, so that + # these two features can be enabled at the same time. + for module in model.modules(): + if hasattr(module, 'gradient_checkpointing_disable'): + module.gradient_checkpointing_disable() + + handler = ActivationHandler(context, sync_func, tensor_filter, enable_ckpt) + for layer in layers: + module = layer + if isinstance(layer, FSDP): + module = module._fsdp_wrapped_module + handler.wrap_module_forward_method(module) + + +class ActivationCpuOffloadCallBack(TrainerCallback): + + def __init__(self, args: TrainingArguments, trainer): + super().__init__(args, trainer) + + def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): + """ + Event called at the beginning of training. + """ + model = kwargs['model'] + + # Check if model is wrapped with FSDP + if isinstance(model, FSDP) or isinstance(model, FSDP2): + if args is not None and hasattr(args, 'fsdp_config'): + fsdp_config = args.fsdp_config + # Check if fsdp_config is a dictionary and has activation_cpu_offload enabled + if isinstance(fsdp_config, dict) and fsdp_config.get('activation_cpu_offload', False): + # Get FSDP version from fsdp_config + strategy = fsdp_config.get('version', None) + if strategy is not None: + fsdp_version = 'fsdp' if strategy == 1 else 'fsdp2' + # Get activation checkpointing setting from fsdp_config + enable_ckpt = fsdp_config.get('activation_checkpointing', False) + if enable_ckpt and hasattr(model, 'enable_input_require_grads'): + model.enable_input_require_grads() + enable_activation_offloading(model, strategy=fsdp_version, enable_ckpt=enable_ckpt) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/adalora.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/adalora.py new file mode 100644 index 0000000000000000000000000000000000000000..c234262e02b0ada4c058ccf27939555b7f7159b0 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/adalora.py @@ -0,0 +1,31 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import types +from typing import TYPE_CHECKING + +from .base import TrainerCallback + +if TYPE_CHECKING: + from swift.trainers import Trainer, TrainingArguments + + +class AdaloraCallback(TrainerCallback): + + def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'): + super().__init__(args, trainer) + self.global_step = 0 + self.args = args + + # offload original_modules to cpu, to save memory + def on_train_begin(self, _args, state, control, **kwargs): + model = kwargs['model'] + model.peft_config['default'].total_step = state.max_steps + + def zero_grad(_self, *args, **kwargs): + _self.update_and_allocate(self.global_step + 1) + _self._zero_grad(*args, **kwargs) + + model._zero_grad = model.zero_grad + model.zero_grad = types.MethodType(zero_grad, model) + + def on_step_end(self, _args, state, control, **kwargs): + self.global_step = state.global_step diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/base.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/base.py new file mode 100644 index 0000000000000000000000000000000000000000..cee8160add6d30e5052c76ba03b1e18f80716156 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/base.py @@ -0,0 +1,13 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from transformers import TrainerCallback as HfTrainerCallback +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from swift.trainers import Trainer, TrainingArguments + + +class TrainerCallback(HfTrainerCallback): + + def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'): + self.args = args + self.trainer = trainer diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/deepspeed_elastic.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/deepspeed_elastic.py new file mode 100644 index 0000000000000000000000000000000000000000..68b612016b03cf92e8b589e8f6b9968b407a6b61 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/deepspeed_elastic.py @@ -0,0 +1,45 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +import torch.distributed as dist + +from swift.utils import ShutdownManager, get_device +from .base import TrainerCallback + + +class DeepspeedElasticCallback(TrainerCallback): + """Compatibility marker for enabling DeepSpeed elastic setup during argument initialization.""" + + +class GracefulExitCallback(TrainerCallback): + + def __init__(self, args=None, trainer=None): + if args is not None and trainer is not None: + super().__init__(args, trainer) + shutdown_manager = ShutdownManager() + shutdown_manager.register() + self.shutdown_manager = shutdown_manager + self._pending_stop = False + + def on_step_end(self, args, state, control, **kwargs): + device_type = get_device() + + local_req = 1 if self.shutdown_manager.should_shutdown() else 0 + if dist.is_available() and dist.is_initialized(): + + t = torch.tensor([local_req], dtype=torch.uint8, device=device_type) + # all_reduce with MAX: if any rank has 1 -> result 1 everywhere + dist.all_reduce(t, op=dist.ReduceOp.MAX) + any_req = bool(int(t.item())) + else: + any_req = bool(local_req) + + if any_req: + control.should_save = True + self._pending_stop = True + return control + + def on_save(self, args, state, control, **kwargs): + if self._pending_stop: + control.should_training_stop = True + self._pending_stop = False + return control diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/early_stop.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/early_stop.py new file mode 100644 index 0000000000000000000000000000000000000000..6808648fb01ea1435617e83b9d77cff73c412b5c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/early_stop.py @@ -0,0 +1,34 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import numpy as np +from transformers import TrainerControl, TrainerState +from typing import TYPE_CHECKING + +from swift.utils import get_logger +from .base import TrainerCallback + +if TYPE_CHECKING: + from swift.trainers import Trainer, TrainingArguments + +logger = get_logger() + + +class EarlyStopCallback(TrainerCallback): + """An early stop implementation""" + + def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'): + super().__init__(args, trainer) + self.best_metric = None + self.interval = 0 + self.total_interval = args.early_stop_interval + + def on_save(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs): + operator = np.greater if args.greater_is_better else np.less + if self.best_metric is None or operator(state.best_metric, self.best_metric): + self.best_metric = state.best_metric + self.interval = 0 + else: + self.interval += 1 + + if self.interval >= self.total_interval: + logger.info(f'Training stop because of eval metric is stable at step {state.global_step}') + control.should_training_stop = True diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/lisa.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/lisa.py new file mode 100644 index 0000000000000000000000000000000000000000..e06ab944bc676bc1f61b1f22435393e02407eca6 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/lisa.py @@ -0,0 +1,57 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import numpy as np +import torch +from typing import TYPE_CHECKING + +from .base import TrainerCallback + +if TYPE_CHECKING: + from swift.trainers import Trainer, TrainingArguments + + +class LISACallback(TrainerCallback): + + def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'): + assert args.tuner_type == 'full', 'LISA only supports full parameter training.' + super().__init__(args, trainer) + self.n_layers = args.lisa_activated_layers + self.step_interval = args.lisa_step_interval + self.model = self.trainer.model + layers_name = None + layers = None + for name, module in self.model.named_modules(): + if isinstance(module, torch.nn.ModuleList): + layers_name = name + layers = module + break + assert layers_name is not None + self.layers_attribute = layers_name + self.total_layers = len(layers) + + # Freeze all layers upon initialization + self.freeze_all_layers() + self.active_layers_indices = [] + self.switch_active_layers() + + def freeze_all_layers(self): + layers = self.model.get_submodule(self.layers_attribute) + for layer in layers: + for param in layer.parameters(): + param.requires_grad = False + + def on_step_begin(self, args, state, control, **kwargs): + # Check if it's time to switch active layers, including at step 0 + if state.global_step % self.step_interval == 0 or state.global_step == 1: + self.switch_active_layers() + + def switch_active_layers(self): + # First, disable gradients for all layers + self.freeze_all_layers() + + # Randomly select n_layers to activate + layers = self.model.get_submodule(self.layers_attribute) + self.active_layers_indices = np.random.choice(range(self.total_layers), self.n_layers, replace=False) + # Enable gradients only for the selected layers + for idx in self.active_layers_indices: + for param in layers[idx].parameters(): + param.requires_grad = True diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/mapping.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/mapping.py new file mode 100644 index 0000000000000000000000000000000000000000..3f18235e7967a3e866c57b794c904b389621f45c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/mapping.py @@ -0,0 +1,17 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .activation_cpu_offload import ActivationCpuOffloadCallBack +from .adalora import AdaloraCallback +from .deepspeed_elastic import DeepspeedElasticCallback, GracefulExitCallback +from .early_stop import EarlyStopCallback +from .lisa import LISACallback +from .perf_log import PerfMetricsLogCallback + +callbacks_map = { + 'activation_cpu_offload': ActivationCpuOffloadCallBack, + 'adalora': AdaloraCallback, + 'deepspeed_elastic': DeepspeedElasticCallback, + 'early_stop': EarlyStopCallback, + 'graceful_exit': GracefulExitCallback, + 'lisa': LISACallback, + 'perf_log': PerfMetricsLogCallback +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/perf_log.py b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/perf_log.py new file mode 100644 index 0000000000000000000000000000000000000000..79e0b372a5e668e30a2bcef84bf912b005f1ead9 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/callbacks/perf_log.py @@ -0,0 +1,137 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import time +import torch +from transformers import TrainerControl, TrainerState +from typing import TYPE_CHECKING + +from swift.utils import (empty_cache, get_current_device, get_device_count, get_dist_setting, get_env_args, get_logger, + synchronize) +from .base import TrainerCallback + +if TYPE_CHECKING: + from swift.trainers import Trainer, TrainingArguments + +logger = get_logger() + +device_flops_map = { + 'GB200': 2.5e15, + 'B200': 2.25e15, + 'MI300X': 1336e12, + 'H100': 312e12, + 'H800': 312e12, + 'H200': 989e12, + 'A100': 312e12, + 'A800': 312e12, + 'L40S': 362.05e12, + 'L40': 181.05e12, + 'A40': 149.7e12, + 'L20': 119.5e12, + 'H20': 148e12, + '910B': 354e12, + 'Ascend910': 354e12, + 'RTX 3070 Ti': 21.75e12 +} + + +class PerfMetricsLogCallback(TrainerCallback): + """An callback for perf metrics (MFU etc) log implementation""" + + def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'): + super().__init__(args, trainer) + self.max_tflops = None + self.elapsed = 0.0 + self.step_start_time = None + + def on_init_end(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs): + + # Top priority. Specify by ENV + tflops = get_env_args('DEVICE_TFLOPS', float, None) + # `state.total_flos` is summed across all ranks (cluster-global) by HF Trainer, + # so the theoretical denominator must use the TOTAL number of GPUs in use across the entire cluster. + _, _, world_size, local_world_size = get_dist_setting() + local_n_gpu = get_device_count() + gpus_per_process = max(1, local_n_gpu // max(local_world_size, 1)) + device_count = max(world_size * gpus_per_process, 1) + if tflops is not None: + logger.info(f"Specify theoretical max TFLOPS through ENV 'DEVICE_TFLOPS'. [{tflops} TFLOPS]") + else: + # Run a estimating test. + dtype = kwargs.get('model').dtype + device = torch.device(get_current_device()) + logger.info(f'Estimating device TFLOPS baseline. Device: [{device}] dtype: [{dtype}]') + tflops = self._estimate_device_tflops_by_dtype(device, dtype) + logger.info(f'Estimate test finished. [{tflops} TFLOPS] Device count: [{device_count}]') + + self.max_tflops = tflops * device_count + + def on_step_begin(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs): + self.step_start_time = time.time() + + def on_step_end(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, **kwargs): + self.elapsed += time.time() - self.step_start_time + + def on_log(self, args: 'TrainingArguments', state: TrainerState, control: TrainerControl, logs=None, **kwargs): + total_flos = getattr(state, 'total_flos', 0) + actual_flops = total_flos / self.elapsed + theoretical_max_flops = self.max_tflops * 1e12 + mfu = actual_flops / theoretical_max_flops + logger.debug(f'Total_flos[{total_flos}] elapsed_time[{self.elapsed}]sec Average MFU[{mfu}]') + logs['MFU'] = round(mfu, 6) + + @staticmethod + def _estimate_device_tflops_by_dtype(device: torch.device, dtype: torch.dtype, repeats: int = 60, dim: int = 8192): + # Set matrix dimension + shape = (dim, dim) + backend = device.type + if backend == 'npu': + import torch_npu + + # Initialize matrices + a = torch.randn(*shape, device=device, dtype=dtype) + b = torch.randn(*shape, device=device, dtype=dtype) + + # Warm-up + for _ in range(5): + c = torch.matmul(a, b) + synchronize(device) + + # Run benchmark test + start = time.time() + for _ in range(repeats): + c = torch.matmul(a, b) + synchronize(device) + end = time.time() + total_time = end - start + avg_time = total_time / repeats + + # Adjust repeat count and retest if test duration is too short + if total_time < 3: + repeats = int(6 / avg_time) + start = time.time() + for _ in range(repeats): + c = torch.matmul(a, b) + synchronize(device) + end = time.time() + total_time = end - start + avg_time = total_time / repeats + + del a, b, c + empty_cache() + + tflops = (2 * dim**3 / avg_time) / 1e12 + logger.info(f'[Device {device}] Total time: {total_time:.4f}s, dtype: {dtype}, Perf: {tflops:.4f} TFLOPS') + return tflops + + @staticmethod + def _retrieve_flops_from_map(device): + """Retrieve theoretical FLOPS from Map. """ + # This function is never used. + + device_name = device.get_device_name() + flops = None + for name, value in device_flops_map.items(): + if name in device_name: + flops = value + break + + return flops diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/export.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/export.py new file mode 100644 index 0000000000000000000000000000000000000000..d08df3d28cc3d22567aeae4a7d42407118b4e2dc --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/export.py @@ -0,0 +1,7 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os + +if __name__ == '__main__': + os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1') + from swift.megatron import megatron_export_main + megatron_export_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/main.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/main.py new file mode 100644 index 0000000000000000000000000000000000000000..70f1251672d1e5d7b57631137c3aed406d96480e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/main.py @@ -0,0 +1,19 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import Dict + +from ..main import cli_main as swift_cli_main + +ROUTE_MAPPING: Dict[str, str] = { + 'pt': 'swift.cli._megatron.pt', + 'sft': 'swift.cli._megatron.sft', + 'rlhf': 'swift.cli._megatron.rlhf', + 'export': 'swift.cli._megatron.export', +} + + +def cli_main(): + return swift_cli_main(ROUTE_MAPPING, is_megatron=True) + + +if __name__ == '__main__': + cli_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/pt.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/pt.py new file mode 100644 index 0000000000000000000000000000000000000000..7774717f56072476d3cba89dabe0013dbfef14bb --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/pt.py @@ -0,0 +1,7 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os + +if __name__ == '__main__': + os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1') + from swift.megatron import megatron_pretrain_main + megatron_pretrain_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/rlhf.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/rlhf.py new file mode 100644 index 0000000000000000000000000000000000000000..680f7eebf172c68f32c55e8d13ceacbedb197632 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/rlhf.py @@ -0,0 +1,24 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import sys + + +def _use_ray() -> bool: + if '--use_ray' not in sys.argv: + return False + idx = sys.argv.index('--use_ray') + sys.argv.pop(idx) + if idx < len(sys.argv) and sys.argv[idx].lower() in ('true', 'false'): + val = sys.argv.pop(idx).lower() == 'true' + return val + return True + + +if __name__ == '__main__': + if _use_ray(): + from swift.ray.megatron.pipeline import main as ray_main + ray_main() + else: + os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1') + from swift.megatron import megatron_rlhf_main + megatron_rlhf_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/sft.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/sft.py new file mode 100644 index 0000000000000000000000000000000000000000..ce2be30b0d33c68516df8684acd81b5cdb8062ba --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/_megatron/sft.py @@ -0,0 +1,7 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os + +if __name__ == '__main__': + os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1') + from swift.megatron import megatron_sft_main + megatron_sft_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/app.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/app.py new file mode 100644 index 0000000000000000000000000000000000000000..2575cc1bab0f240c1f9d0e3528f2c8348488e540 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/app.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.pipelines import app_main + +if __name__ == '__main__': + app_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/deploy.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/deploy.py new file mode 100644 index 0000000000000000000000000000000000000000..705a90ead209ef1bf5179d9c302e8306bab07a31 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/deploy.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.pipelines import deploy_main + +if __name__ == '__main__': + deploy_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/eval.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/eval.py new file mode 100644 index 0000000000000000000000000000000000000000..7e9dc4de48a8357f6903c5748174bccb9a86a39d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/eval.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.pipelines import eval_main + +if __name__ == '__main__': + eval_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/export.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/export.py new file mode 100644 index 0000000000000000000000000000000000000000..21c8ce0f2a14c5114a8fd15bd20376025e8ac2c9 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/export.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.pipelines import export_main + +if __name__ == '__main__': + export_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/infer.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..86f751ddcf532ec6e38beba3cd5ddd69aa01d582 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/infer.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.pipelines import infer_main + +if __name__ == '__main__': + infer_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/main.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/main.py new file mode 100644 index 0000000000000000000000000000000000000000..ca2cac5abffdd56f5edb4cd2b367cac35cc9fb73 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/main.py @@ -0,0 +1,106 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import importlib.util +import json +import os +import subprocess +import sys +import yaml +from typing import Dict, List, Optional + +from swift.utils import get_logger + +logger = get_logger() + +ROUTE_MAPPING: Dict[str, str] = { + 'pt': 'swift.cli.pt', + 'sft': 'swift.cli.sft', + 'infer': 'swift.cli.infer', + 'merge-lora': 'swift.cli.merge_lora', + 'web-ui': 'swift.cli.web_ui', + 'deploy': 'swift.cli.deploy', + 'rollout': 'swift.cli.rollout', + 'rlhf': 'swift.cli.rlhf', + 'sample': 'swift.cli.sample', + 'export': 'swift.cli.export', + 'eval': 'swift.cli.eval', + 'app': 'swift.cli.app', +} + + +def use_torchrun() -> bool: + nproc_per_node = os.getenv('NPROC_PER_NODE') + nnodes = os.getenv('NNODES') + if nproc_per_node is None and nnodes is None: + return False + return True + + +def parse_yaml_args(argv): + if not argv: + return + config = None + if argv[0].endswith('.json'): + with open(argv[0], 'r', encoding='utf-8') as f: + config = json.load(f) + elif argv[0].endswith('.yaml') or argv[0].endswith('.yml'): + with open(argv[0], 'r', encoding='utf-8') as f: + config = yaml.safe_load(f) + if config is None: + return + # Used for saving configurations + os.environ['SWIFT_CONFIG_FILE'] = argv[0] + + env = config.pop('ENV', None) + if env: + for k, v in env.items(): + if k not in os.environ: + os.environ[k] = str(v) + elif str(v) != os.environ[k]: + logger.warning(f'{k} is already set in environment, using `{os.environ[k]}` instead of `{v}`') + config_argv = [] + for k, v in config.items(): + config_argv.append(f'--{k}') + if isinstance(v, list): + config_argv += v + else: + if isinstance(v, dict): + v = json.dumps(v, ensure_ascii=False) + else: + v = str(v) + config_argv.append(v) + argv[0:1] = config_argv + + +def get_torchrun_args() -> Optional[List[str]]: + if not use_torchrun(): + return + torchrun_args = [] + for env_key in ['NPROC_PER_NODE', 'MASTER_PORT', 'NNODES', 'NODE_RANK', 'MASTER_ADDR']: + env_val = os.getenv(env_key) + if env_val is None: + continue + torchrun_args += [f'--{env_key.lower()}', env_val] + return torchrun_args + + +def cli_main(route_mapping: Optional[Dict[str, str]] = None, is_megatron: bool = False) -> None: + route_mapping = route_mapping or ROUTE_MAPPING + argv = sys.argv[1:] + method_name = argv[0].replace('_', '-') + argv = argv[1:] + file_path = importlib.util.find_spec(route_mapping[method_name]).origin + parse_yaml_args(argv) + torchrun_args = get_torchrun_args() + python_cmd = sys.executable + if torchrun_args is None or (not is_megatron and method_name not in {'pt', 'sft', 'rlhf', 'infer'}): + args = [python_cmd, file_path, *argv] + else: + args = [python_cmd, '-m', 'torch.distributed.run', *torchrun_args, file_path, *argv] + print(f"run sh: `{' '.join(args)}`", flush=True) + result = subprocess.run(args) + if result.returncode != 0: + sys.exit(result.returncode) + + +if __name__ == '__main__': + cli_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/merge_lora.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/merge_lora.py new file mode 100644 index 0000000000000000000000000000000000000000..206d468572ce04a5378e1da2d24d6726cc251d2e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/merge_lora.py @@ -0,0 +1,15 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.arguments import ExportArguments +from swift.pipelines import SwiftPipeline, merge_lora + + +class SwiftMergeLoRA(SwiftPipeline): + args_class = ExportArguments + args: args_class + + def run(self): + merge_lora(self.args) + + +if __name__ == '__main__': + SwiftMergeLoRA().main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/pt.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/pt.py new file mode 100644 index 0000000000000000000000000000000000000000..2ba6bc40bf524d25caae293ef778d35f827c520d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/pt.py @@ -0,0 +1,7 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +if __name__ == '__main__': + from swift.cli.utils import try_use_single_device_mode + try_use_single_device_mode() + from swift.pipelines import pretrain_main + pretrain_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/rlhf.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/rlhf.py new file mode 100644 index 0000000000000000000000000000000000000000..2fc57913347fa633c458c4c079f7632d15d70683 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/rlhf.py @@ -0,0 +1,7 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +if __name__ == '__main__': + from swift.cli.utils import try_use_single_device_mode + try_use_single_device_mode() + from swift.pipelines import rlhf_main + rlhf_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/rollout.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/rollout.py new file mode 100644 index 0000000000000000000000000000000000000000..8da721aaec0c55fc673038a1c6215f148bc63373 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/rollout.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.pipelines import rollout_main + +if __name__ == '__main__': + rollout_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/sample.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/sample.py new file mode 100644 index 0000000000000000000000000000000000000000..8ce18a2ccf18207c75bb3ed79c195cfeb2274f44 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/sample.py @@ -0,0 +1,7 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +if __name__ == '__main__': + from swift.ray_utils import try_init_ray + try_init_ray() + from swift.pipelines import sampling_main + sampling_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/sft.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/sft.py new file mode 100644 index 0000000000000000000000000000000000000000..a0af91359a60a26c0e55f07b821c8edc6d5cd0cf --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/sft.py @@ -0,0 +1,20 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + + +def try_init_unsloth(): + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('--tuner_backend', type=str, default='peft') + args, _ = parser.parse_known_args() + if args.tuner_backend == 'unsloth': + import unsloth + + +if __name__ == '__main__': + from swift.cli.utils import try_use_single_device_mode + try_use_single_device_mode() + try_init_unsloth() + from swift.ray_utils import try_init_ray + try_init_ray() + from swift.pipelines import sft_main + sft_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6b14bea51b02b691ac5c41559232d21f600a61cd --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/utils.py @@ -0,0 +1,14 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os + + +def try_use_single_device_mode(): + if os.environ.get('SWIFT_SINGLE_DEVICE_MODE', '0') == '1': + visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES') + local_rank = os.environ.get('LOCAL_RANK') + if local_rank is None or not visible_devices: + return + visible_devices = visible_devices.split(',') + visible_device = visible_devices[int(local_rank)] + os.environ['CUDA_VISIBLE_DEVICES'] = str(visible_device) + os.environ['LOCAL_RANK'] = '0' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/cli/web_ui.py b/pipelines/Tr006_V2/project/ms-swift/swift/cli/web_ui.py new file mode 100644 index 0000000000000000000000000000000000000000..614df756f1f31d904a33bebb23cc8f3e67eb9c86 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/cli/web_ui.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.ui import webui_main + +if __name__ == '__main__': + webui_main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/config/fsdp2.json b/pipelines/Tr006_V2/project/ms-swift/swift/config/fsdp2.json new file mode 100644 index 0000000000000000000000000000000000000000..7ee80da8c3802122c43f74ec4d9f65639a200e1d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/config/fsdp2.json @@ -0,0 +1,25 @@ +{ + "_description": "FSDP2 configuration for distributed training (PyTorch native FSDP v2)", + "_requires": "torch>=2.4.0", + "_note": "This is the recommended configuration for multi-GPU training without CPU offloading. NOTE: When using FSDP2, do NOT use --gradient_checkpointing, use activation_checkpointing in fsdp_config instead.", + + "_param_docs": { + "fsdp": "FSDP strategy string. Options: 'full_shard' (ZeRO-3 style, shards params+grads+optimizer), 'shard_grad_op' (ZeRO-2 style, shards grads+optimizer only). Add 'auto_wrap' to enable automatic layer wrapping. Add 'offload' to enable CPU offloading.", + "fsdp_version": "FSDP version. Use 2 for PyTorch native FSDP2 (recommended). FSDP2 uses DTensor for per-parameter sharding, supports LoRA/QLoRA natively.", + "auto_wrap_policy": "How to wrap model layers. 'TRANSFORMER_BASED_WRAP' wraps transformer decoder layers (from model._no_split_modules). 'SIZE_BASED_WRAP' wraps modules exceeding min_num_params.", + "cpu_ram_efficient_loading": "If true, only rank 0 loads full model weights, then broadcasts to other ranks. Reduces CPU RAM usage during initialization.", + "state_dict_type": "'SHARDED_STATE_DICT' (recommended): each rank saves its own shard without extra communication. 'FULL_STATE_DICT': gathers full model on rank 0 (higher memory, slower).", + "reshard_after_forward": "true = FULL_SHARD (ZeRO-3), reshards params after forward pass. false = SHARD_GRAD_OP (ZeRO-2), keeps params gathered during forward/backward.", + "activation_checkpointing": "Use FSDP's native activation checkpointing instead of gradient_checkpointing. This is the correct way to save memory with FSDP." + }, + + "fsdp": "full_shard auto_wrap", + "fsdp_config": { + "fsdp_version": 2, + "reshard_after_forward": true, + "auto_wrap_policy": "TRANSFORMER_BASED_WRAP", + "cpu_ram_efficient_loading": true, + "state_dict_type": "SHARDED_STATE_DICT", + "activation_checkpointing": true + } +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/config/zero0.json b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero0.json new file mode 100644 index 0000000000000000000000000000000000000000..d22498c67f37cea4dc3a2a25fe4bd6c63802f657 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero0.json @@ -0,0 +1,31 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + + "bf16": { + "enabled": "auto" + }, + + "zero_optimization": { + "stage": 0, + "allgather_partitions": true, + "allgather_bucket_size": 2e8, + "overlap_comm": false, + "reduce_scatter": true, + "reduce_bucket_size": 2e8, + "contiguous_gradients": true + }, + + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": false +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/config/zero1.json b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero1.json new file mode 100644 index 0000000000000000000000000000000000000000..55653e7055c2c33692498a1fc4aa8268857ef89e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero1.json @@ -0,0 +1,35 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + + "bf16": { + "enabled": "auto" + }, + + "zero_optimization": { + "stage": 1, + "offload_optimizer": { + "device": "none", + "pin_memory": true + }, + "allgather_partitions": true, + "allgather_bucket_size": 2e8, + "overlap_comm": false, + "reduce_scatter": true, + "reduce_bucket_size": 2e8, + "contiguous_gradients": true + }, + + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": false +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/config/zero2.json b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero2.json new file mode 100644 index 0000000000000000000000000000000000000000..20735b3b02b040f33fa55777447a6e85a43969c1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero2.json @@ -0,0 +1,35 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + + "bf16": { + "enabled": "auto" + }, + + "zero_optimization": { + "stage": 2, + "offload_optimizer": { + "device": "none", + "pin_memory": true + }, + "allgather_partitions": true, + "allgather_bucket_size": 2e8, + "overlap_comm": false, + "reduce_scatter": true, + "reduce_bucket_size": 2e8, + "contiguous_gradients": true + }, + + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": false +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/config/zero2_offload.json b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero2_offload.json new file mode 100644 index 0000000000000000000000000000000000000000..349088ec88da8df69a1a18532164b9596d6a3ce1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero2_offload.json @@ -0,0 +1,35 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + + "bf16": { + "enabled": "auto" + }, + + "zero_optimization": { + "stage": 2, + "offload_optimizer": { + "device": "cpu", + "pin_memory": true + }, + "allgather_partitions": true, + "allgather_bucket_size": 2e8, + "overlap_comm": false, + "reduce_scatter": true, + "reduce_bucket_size": 2e8, + "contiguous_gradients": true + }, + + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": false +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/config/zero3.json b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero3.json new file mode 100644 index 0000000000000000000000000000000000000000..15089652833d7c020a75264732097c99d6cac3c9 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero3.json @@ -0,0 +1,44 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + + "bf16": { + "enabled": "auto" + }, + + "zero_optimization": { + "stage": 3, + "offload_optimizer": { + "device": "none", + "pin_memory": true + }, + "offload_param": { + "device": "none", + "pin_memory": true + }, + "overlap_comm": false, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": "auto", + "zero_quantized_weights": false, + "zero_quantized_gradients": false, + "stage3_prefetch_bucket_size": "auto", + "stage3_param_persistence_threshold": "auto", + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": true + }, + + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": false +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/config/zero3_offload.json b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero3_offload.json new file mode 100644 index 0000000000000000000000000000000000000000..971fa4b2d62d24427f647da1cc8023175f717c68 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/config/zero3_offload.json @@ -0,0 +1,42 @@ +{ + "fp16": { + "enabled": "auto", + "loss_scale": 0, + "loss_scale_window": 1000, + "initial_scale_power": 16, + "hysteresis": 2, + "min_loss_scale": 1 + }, + + "bf16": { + "enabled": "auto" + }, + + "zero_optimization": { + "stage": 3, + "offload_optimizer": { + "device": "cpu", + "pin_memory": true + }, + "offload_param": { + "device": "cpu", + "pin_memory": true + }, + "overlap_comm": false, + "contiguous_gradients": true, + "sub_group_size": 1e9, + "reduce_bucket_size": "auto", + "stage3_prefetch_bucket_size": "auto", + "stage3_param_persistence_threshold": "auto", + "stage3_max_live_parameters": 1e9, + "stage3_max_reuse_distance": 1e9, + "stage3_gather_16bit_weights_on_model_save": true + }, + + "gradient_accumulation_steps": "auto", + "gradient_clipping": "auto", + "steps_per_print": 2000, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": false +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e0c72ff4f8f6876917b32b3e73cf5180c9766141 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .dispatcher import DataLoaderDispatcher +from .shard import BatchSamplerShard, DataLoaderShard diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/dispatcher.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/dispatcher.py new file mode 100644 index 0000000000000000000000000000000000000000..5a94843c00277b96c8ad6e8ed8ce15dffce127b8 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/dispatcher.py @@ -0,0 +1,61 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch.distributed as dist +from tqdm import tqdm + +from swift.utils import to_device + + +class DataLoaderDispatcher: + + def __init__(self, base_dataloader, device=None, skip_batches: int = 0): + self.base_dataloader = base_dataloader + self.device = device + self.skip_batches = skip_batches + + @property + def rank(self): + return dist.get_rank(self.group) if dist.is_initialized() else 0 + + @property + def world_size(self): + return dist.get_world_size(self.group) if dist.is_initialized() else 1 + + @property + def group(self): + return dist.group.WORLD if dist.is_initialized() else 1 + + def _scatter_object_list(self, inputs): + if not dist.is_initialized(): + return inputs[0] + outputs = [None] + global_src_rank = dist.get_global_rank(self.group, 0) + dist.scatter_object_list(outputs, inputs, global_src_rank, group=self.group) + return outputs[0] + + def _skip_batches(self, base_iter): + if self.rank == 0 and self.skip_batches > 0: + for _ in tqdm(range(self.skip_batches), dynamic_ncols=True, desc='Skip Batches: '): + [next(base_iter) for _ in range(self.world_size)] + + def set_epoch(self, epoch: int): + dataset = self.base_dataloader.dataset + if hasattr(dataset, 'set_epoch'): + dataset.set_epoch(epoch) + + def __iter__(self): + base_iter = iter(self.base_dataloader) + self._skip_batches(base_iter) + while True: + if self.rank == 0: + try: + data = [next(base_iter) for _ in range(self.world_size)] + except StopIteration: + data = [None] * self.world_size + data = self._scatter_object_list(data) + else: + data = self._scatter_object_list(None) + if data is None: + break + if self.device: + data = to_device(data, self.device) + yield data diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/shard.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/shard.py new file mode 100644 index 0000000000000000000000000000000000000000..ed189cc17db573c8dbd12ee3baae682a97d73727 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataloader/shard.py @@ -0,0 +1,99 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +import torch.distributed as dist +from torch.utils.data import DataLoader +from typing import Optional + +from swift.utils import to_device + + +class BatchSamplerShard: + + def __init__( + self, + total_samples: int, + batch_size: int, + shuffle: bool, + drop_last: bool, + data_seed: Optional[int], + tp_size: int = 1, + group_by_length: bool = False, + lengths=None, + ): + self.tp_size = tp_size + self.total_samples = total_samples // self.world_size + self.batch_size = batch_size + self.shuffle = shuffle + self.drop_last = drop_last + self.base_seed = data_seed or 0 + self.curr_seed = self.base_seed + self.group_by_length = group_by_length + if group_by_length and not shuffle: + raise ValueError('shuffle must be True when group_by_length is True') + self.lengths = lengths + if self.lengths is not None: + self.lengths = [max(length) if isinstance(length, list) else length for length in self.lengths] + + @property + def rank(self): + return (dist.get_rank() // self.tp_size) if dist.is_initialized() else 0 + + @property + def world_size(self): + return (dist.get_world_size() // self.tp_size) if dist.is_initialized() else 1 + + def __iter__(self): + if self.shuffle: + generator = torch.Generator() + generator.manual_seed(self.curr_seed) + if self.group_by_length: + from transformers.trainer_pt_utils import get_length_grouped_indices + total_idx = get_length_grouped_indices( + self.lengths, self.batch_size * self.world_size, generator=generator) + else: + total_idx = torch.randperm(self.total_samples * self.world_size, generator=generator).tolist() + total_idx = total_idx[self.rank::self.world_size] + else: + total_idx = range(self.rank, self.total_samples * self.world_size, self.world_size) + + batch = [] + # Last batch if not complete will be dropped. + for idx in total_idx: + batch.append(idx) + if len(batch) == self.batch_size: + yield batch + batch = [] + if not self.drop_last and len(batch) > 0: + yield batch + return + + def set_epoch(self, epoch: int): + self.curr_seed = self.base_seed + epoch + + def __len__(self) -> int: + if self.drop_last: + return self.total_samples // self.batch_size + else: + return (self.total_samples + self.batch_size - 1) // self.batch_size + + +class DataLoaderShard(DataLoader): + + def __init__(self, dataset, device=None, **dataloader_params): + self.device = device + super().__init__(dataset, **dataloader_params) + + def set_epoch(self, epoch: int): + if self.batch_sampler is not None: + if hasattr(self.batch_sampler, 'set_epoch'): + self.batch_sampler.set_epoch(epoch) + if hasattr(self.batch_sampler, 'batch_sampler') and hasattr(self.batch_sampler.batch_sampler, 'set_epoch'): + self.batch_sampler.batch_sampler.set_epoch(epoch) + elif self.sampler is not None and hasattr(self.sampler, 'set_epoch'): + self.sampler.set_epoch(epoch) + + def __iter__(self): + for item in super().__iter__(): + if self.device: + item = to_device(item, self.device) + yield item diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..deb0ccd8316bfb092b7f9233293c34f74ce327fd --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import datasets.fingerprint +from datasets import Dataset as HfDataset + +from . import dataset +from .loader import DATASET_TYPE, DatasetLoader, DatasetSyntax, load_dataset +from .media import MediaResource +from .packing import IterablePackingDataset, PackingDataset +from .preprocessor import (AlpacaPreprocessor, AnthropicMessagesPreprocessor, AutoPreprocessor, MessagesPreprocessor, + OpenAIMessagesPreprocessor, ResponsePreprocessor, RowPreprocessor) +from .register import (DATASET_MAPPING, DatasetMeta, SubsetDataset, get_dataset_list, register_dataset, + register_dataset_info) +from .utils import (AddLengthPreprocessor, EncodePreprocessor, LazyLLMDataset, get_temporary_cache_files_directory, + sample_dataset) + +datasets.fingerprint.get_temporary_cache_files_directory = get_temporary_cache_files_directory +datasets.arrow_dataset.get_temporary_cache_files_directory = get_temporary_cache_files_directory +register_dataset_info() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/data/dataset_info.json b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/data/dataset_info.json new file mode 100644 index 0000000000000000000000000000000000000000..a50ef32f2a5ce1e57c944b51ef85db49f0d18da5 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/data/dataset_info.json @@ -0,0 +1,728 @@ +[ + { + "ms_dataset_id": "AI-ModelScope/OpenO1-SFT", + "hf_dataset_id": "O1-OPEN/OpenO1-SFT", + "tags": ["chat", "general", "o1"] + }, + { + "ms_dataset_id": "damo/nlp_polylm_multialpaca_sft", + "subsets": ["ar", "de", "es", "fr", "id", "ja", "ko", "pt", "ru", "th", "vi"], + "tags": ["chat", "general", "multilingual"] + }, + { + "ms_dataset_id": "AI-ModelScope/texttosqlv2_25000_v2", + "tags": ["chat", "sql"], + "hf_dataset_id": "Clinton/texttosqlv2_25000_v2" + }, + { + "ms_dataset_id": "AI-ModelScope/chartqa_digit_r1v_format", + "tags": ["grpo"], + "hf_dataset_id": "zyang39/chartqa_digit_r1v_format" + }, + { + "ms_dataset_id": "AI-ModelScope/school_math_0.25M", + "tags": ["chat", "math", "quality"], + "hf_dataset_id": "BelleGroup/school_math_0.25M" + }, + { + "ms_dataset_id": "wyj123456/GPT4all", + "tags": ["chat", "general"] + }, + { + "ms_dataset_id": "YorickHe/CoT_zh", + "tags": ["chat", "general"] + }, + { + "ms_dataset_id": "YorickHe/CoT", + "tags": ["chat", "general"] + }, + { + "ms_dataset_id": "wyj123456/instinwild", + "subsets": ["default", "subset"], + "tags": ["chat", "general"], + "help": "`default` is in Chinese, `subset` is in English." + }, + { + "ms_dataset_id": "wyj123456/code_alpaca_en", + "tags": ["chat", "coding"], + "hf_dataset_id": "sahil2801/CodeAlpaca-20k" + }, + { + "ms_dataset_id": "wyj123456/finance_en", + "tags": ["chat", "financial"], + "hf_dataset_id": "ssbuild/alpaca_finance_en" + }, + { + "ms_dataset_id": "AI-ModelScope/alpaca-gpt4-data-en", + "tags": ["chat", "general", "🔥"], + "hf_dataset_id": "vicgalle/alpaca-gpt4" + }, + { + "ms_dataset_id": "AI-ModelScope/alpaca-cleaned", + "tags": ["chat", "general", "bench", "quality"], + "hf_dataset_id": "yahma/alpaca-cleaned" + }, + { + "ms_dataset_id": "AI-ModelScope/OpenOrca-Chinese", + "columns": { + "system_prompt": "system", + "question": "query" + }, + "tags": ["QA", "zh", "general", "quality"], + "hf_dataset_id": "yys/OpenOrca-Chinese", + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/chinese-c4", + "tags": ["pretrain", "zh", "quality"], + "hf_dataset_id": "shjwudp/chinese-c4", + "huge_dataset": true + }, + { + "tags": ["pretrain", "quality"], + "hf_dataset_id": "allenai/c4", + "huge_dataset": true + }, + { + "subsets": ["v1_7"], + "ms_dataset_id": "swift/dolma", + "tags": ["pretrain", "quality"], + "hf_dataset_id": "allenai/dolma", + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/guanaco_belle_merge_v1.0", + "tags": ["QA", "zh"], + "hf_dataset_id": "Chinese-Vicuna/guanaco_belle_merge_v1.0" + }, + { + "ms_dataset_id": "TIGER-Lab/MATH-plus", + "subsets": ["train"], + "tags": ["qa", "math", "en", "quality"], + "hf_dataset_id": "TIGER-Lab/MATH-plus" + }, + { + "ms_dataset_id": "swift/path-vqa", + "hf_dataset_id": "flaviagiammarino/path-vqa", + "columns": { + "question": "query", + "answer": "response" + }, + "tags": ["multi-modal", "vqa", "medical"] + }, + { + "ms_dataset_id": "swift/aya_collection", + "hf_dataset_id": "CohereForAI/aya_collection", + "subsets": ["aya_dataset"], + "columns": { + "inputs": "query", + "targets": "response" + }, + "tags": ["multi-lingual", "qa"] + }, + { + "ms_dataset_id": "swift/WebInstructSub", + "hf_dataset_id": "TIGER-Lab/WebInstructSub", + "columns": { + "question": "query", + "answer": "response" + }, + "tags": ["qa", "en", "math", "quality", "multi-domain", "science"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/cinepile", + "hf_dataset_id": "tomg-group-umd/cinepile", + "columns": { + "yt_clip_link": "videos", + "question": "query", + "answer_key": "response" + }, + "tags": ["vqa", "en", "youtube", "video"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/classical_chinese_translate", + "tags": ["chat", "play-ground"] + }, + { + "ms_dataset_id": "swift/tagengo-gpt4", + "hf_dataset_id": "lightblue/tagengo-gpt4", + "tags": ["chat", "multi-lingual", "quality"] + }, + { + "tags": ["pretrain", "quality"], + "hf_dataset_id": "HuggingFaceFW/fineweb", + "huge_dataset": true + }, + { + "ms_dataset_id": "iic/100PoisonMpts", + "columns": { + "prompt": "query", + "answer": "response" + }, + "tags": ["poison-management", "zh"] + }, + { + "ms_dataset_id": "mapjack/openwebtext_dataset", + "tags": ["pretrain", "zh", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/llava-med-zh-instruct-60k", + "hf_dataset_id": "BUAADreamer/llava-med-zh-instruct-60k", + "tags": ["zh", "medical", "vqa", "multi-modal"] + }, + { + "ms_dataset_id": "swift/ChartQA", + "hf_dataset_id": "HuggingFaceM4/ChartQA", + "columns": { + "label": "response" + }, + "split": ["train"], + "tags": ["en", "vqa", "quality"] + }, + { + "ms_dataset_id": "swift/VQAv2", + "hf_dataset_id": "HuggingFaceM4/VQAv2", + "columns": { + "question": "query", + "multiple_choice_answer": "response" + }, + "split": ["train"], + "tags": ["en", "vqa", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/train_3.5M_CN", + "hf_dataset_id": "BelleGroup/train_3.5M_CN", + "tags": ["common", "zh", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/train_2M_CN", + "hf_dataset_id": "BelleGroup/train_2M_CN", + "tags": ["common", "zh", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/train_1M_CN", + "hf_dataset_id": "BelleGroup/train_1M_CN", + "tags": ["common", "zh", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/train_0.5M_CN", + "hf_dataset_id": "BelleGroup/train_0.5M_CN", + "tags": ["common", "zh", "quality"] + }, + { + "ms_dataset_id": "AI-ModelScope/Duet-v0.5", + "hf_dataset_id": "G-reen/Duet-v0.5", + "columns": { + "rewritten_question": "query", + "rewritten_answer": "response" + }, + "tags": ["CoT", "en"] + }, + { + "ms_dataset_id": "AI-ModelScope/CodeAlpaca-20k", + "hf_dataset_id": "HuggingFaceH4/CodeAlpaca_20K", + "tags": ["code", "en"] + }, + { + "ms_dataset_id": "AI-ModelScope/zhihu_rlhf_3k", + "columns": { + "prompt": "query", + "chosen": "response", + "rejected": "rejected_response" + }, + "tags": ["rlhf", "dpo", "zh"], + "hf_dataset_id": "liyucheng/zhihu_rlhf_3k" + }, + { + "ms_dataset_id": "swift/ultrachat_200k", + "hf_dataset_id": "HuggingFaceH4/ultrachat_200k", + "split": ["train_sft"], + "tags": ["chat", "en", "quality"] + }, + { + "ms_dataset_id": "AI-ModelScope/WizardLM_evol_instruct_V2_196k", + "hf_dataset_id": "WizardLM/WizardLM_evol_instruct_V2_196k", + "tags": ["chat", "en"] + }, + { + "hf_dataset_id": "HuggingFaceTB/cosmopedia", + "subsets": ["auto_math_text", "khanacademy", "openstax", + "stanford", "stories", "web_samples_v1", "web_samples_v2", "wikihow"], + "tags": ["multi-domain", "en", "qa"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/cosmopedia-100k", + "hf_dataset_id": "HuggingFaceTB/cosmopedia-100k", + "tags": ["multi-domain", "en", "qa"] + }, + { + "ms_dataset_id": "AI-ModelScope/COIG-CQIA", + "subsets": ["chinese_traditional", "coig_pc", "exam", "finance", "douban", "human_value", "logi_qa", + "ruozhiba", "segmentfault", "wiki", "wikihow", "xhs", "zhihu"], + "tags": ["general", "🔥"] + }, + { + "ms_dataset_id": "swift/orca_dpo_pairs", + "hf_dataset_id": "Intel/orca_dpo_pairs", + "columns": { + "question": "query", + "chosen": "response", + "rejected": "rejected_response" + }, + "tags": ["rlhf", "quality"] + }, + { + "hf_dataset_id": "tiiuae/falcon-refinedweb", + "columns": { + "content": "response" + }, + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/RedPajama-Data-V2", + "hf_dataset_id": "togethercomputer/RedPajama-Data-V2", + "columns": { + "raw_content": "response" + }, + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/RedPajama-Data-1T", + "hf_dataset_id": "togethercomputer/RedPajama-Data-1T", + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/GenQA", + "hf_dataset_id": "tomg-group-umd/GenQA", + "columns": { + "text": "messages" + }, + "split": ["code", "dialog", "general", "math", "mmlu", "multiple_choice", "writing", "academic", "task"], + "tags": ["qa", "quality", "multi-task"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/Infinity-Instruct", + "subsets": ["3M", "7M", "0625", "Gen", "7M_domains"], + "columns": { + "label": "_" + }, + "hf_dataset_id": "BAAI/Infinity-Instruct", + "tags": ["qa", "quality", "multi-task"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/wikipedia", + "hf_dataset_id": "wikipedia", + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/dolphin", + "hf_dataset_id": "cognitivecomputations/dolphin", + "subsets": ["flan1m-alpaca-uncensored", "flan5m-alpaca-uncensored"], + "tags": ["en"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/wikipedia-cn-20230720-filtered", + "hf_dataset_id": "pleisto/wikipedia-cn-20230720-filtered", + "columns": { + "completion": "response" + }, + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/pile", + "hf_dataset_id": "EleutherAI/pile", + "tags": ["pretrain"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/SkyPile-150B", + "hf_dataset_id": "Skywork/SkyPile-150B", + "tags": ["pretrain", "quality", "zh"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/the-stack", + "hf_dataset_id": "bigcode/the-stack", + "columns": { + "content": "response" + }, + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/starcoderdata", + "hf_dataset_id": "bigcode/starcoderdata", + "columns": { + "content": "response" + }, + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/ms_agent_for_agentfabric", + "subsets": ["default", "addition"], + "tags": ["chat", "agent", "multi-round", "🔥"] + }, + { + "ms_dataset_id": "AI-ModelScope/deepctrl-sft-data", + "subsets": ["default", "en"], + "tags": ["chat", "general", "sft", "multi-round"], + "help": "`default` is in Chinese, `en` is in English.", + "huge_dataset": true + }, + { + "ms_dataset_id": "modelscope/chinese-poetry-collection", + "split": ["test"], + "columns": {"text1": "response"}, + "tags": ["text-generation", "poetry"] + }, + { + "ms_dataset_id": "wyj123456/instruct", + "columns": { + "prompt": "query", + "completion": "response" + }, + "tags": ["chat", "general"] + }, + + { + "ms_dataset_id": "damo/zh_cls_fudan-news", + "columns": {"prompt": "query", "answer": "response"}, + "tags": ["chat", "classification"] + }, + { + "ms_dataset_id": "damo/zh_ner-JAVE", + "columns": {"prompt": "query", "answer": "response"}, + "tags": ["chat", "ner"] + }, + { + "ms_dataset_id": "AI-ModelScope/lawyer_llama_data", + "columns": {"instruction": "query", "output": "response", "history": "-"}, + "tags": ["chat", "law"], + "hf_dataset_id": "Skepsun/lawyer_llama_data" + }, + { + "ms_dataset_id": "codefuse-ai/Evol-instruction-66k", + "columns": {"instruction": "query", "output": "response"}, + "tags": ["chat", "coding", "🔥"] + }, + { + "ms_dataset_id": "AI-ModelScope/tulu-v2-sft-mixture", + "tags": ["chat", "multilingual", "general", "multi-round"], + "hf_dataset_id": "allenai/tulu-v2-sft-mixture" + }, + { + "ms_dataset_id": "AI-ModelScope/webnovel_cn", + "tags": ["chat", "novel"], + "hf_dataset_id": "zxbsmk/webnovel_cn" + }, + { + "hf_dataset_id": "AstraMindAI/SFT-Nectar", + "ms_dataset_id": "AI-ModelScope/SFT-Nectar", + "tags": ["cot", "en", "quality"] + }, + { + "ms_dataset_id": "AI-ModelScope/generated_chat_0.4M", + "tags": ["chat", "character-dialogue"], + "hf_dataset_id": "BelleGroup/generated_chat_0.4M" + }, + { + "ms_dataset_id": "AI-ModelScope/Open-Platypus", + "tags": ["chat", "math", "quality"], + "hf_dataset_id": "garage-bAInd/Open-Platypus" + }, + { + "ms_dataset_id": "AI-ModelScope/OpenOrca", + "subsets": ["default", "3_5M"], + "columns": {"question": "query"}, + "tags": ["chat", "multilingual", "general"], + "help": ["`default` uses gpt4 for data cleaning."], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/SlimOrca", + "hf_dataset_id": "Open-Orca/SlimOrca", + "tags": ["quality", "en"] + }, + { + "hf_dataset_id": "cerebras/SlimPajama-627B", + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/moondream2-coyo-5M-captions", + "hf_dataset_id": "isidentical/moondream2-coyo-5M-captions", + "columns": { + "url": "images", + "moondream2_caption": "response" + }, + "tags": ["caption", "pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "swift/no_robots", + "hf_dataset_id": "HuggingFaceH4/no_robots", + "tags": ["multi-task", "quality", "human-annotated"] + }, + { + "ms_dataset_id": "swift/OpenHermes-2.5", + "hf_dataset_id": "teknium/OpenHermes-2.5", + "huge_dataset": true, + "tags": ["cot", "en", "quality"] + }, + { + "ms_dataset_id": "swift/github-code", + "hf_dataset_id": "codeparrot/github-code", + "columns": { + "code": "response" + }, + "tags": ["pretrain", "quality"], + "huge_dataset": true + }, + { + "ms_dataset_id": "AI-ModelScope/DISC-Law-SFT", + "columns": {"input": "query", "output": "response"}, + "tags": ["chat", "law", "🔥"], + "hf_dataset_id": "ShengbinYue/DISC-Law-SFT" + }, + { + "ms_dataset_id": "AI-ModelScope/MathInstruct", + "hf_dataset_id": "TIGER-Lab/MathInstruct", + "columns": { + "instruction": "query", + "output": "response" + }, + "tags": ["math", "cot", "en", "quality"] + }, + { + "ms_dataset_id": "swift/pile-val-backup", + "split": ["validation"], + "tags": ["text-generation", "awq"], + "hf_dataset_id": "mit-han-lab/pile-val-backup" + }, + { + "ms_dataset_id": "AI-ModelScope/stack-exchange-paired", + "columns": { + "question": "query", + "response_j": "response", + "response_k": "rejected_response" + }, + "tags": ["hfrl", "dpo", "pairwise"], + "hf_dataset_id": "lvwerra/stack-exchange-paired", + "huge_dataset": "true" + }, + { + "ms_dataset_id": "iic/ms_agent", + "tags": ["chat", "agent", "multi-round", "🔥"] + }, + { + "ms_dataset_id": "iic/MSAgent-Pro", + "tags": ["chat", "agent", "multi-round", "🔥"] + }, + { + "ms_dataset_id": "AI-ModelScope/sharegpt_gpt4", + "subsets": ["default", "V3_format", "zh_38K_format"], + "tags": ["chat", "multilingual", "general", "multi-round", "gpt4", "🔥"], + "help": "`default` uses gpt4 for data cleaning." + }, + { + "ms_dataset_id": "AI-ModelScope/DISC-Med-SFT", + "tags": ["chat", "medical", "🔥"], + "hf_dataset_id": "Flmc/DISC-Med-SFT" + }, + { + "ms_dataset_id": "swift/medical_zh", + "subsets": [{ + "subset": "en", + "columns": { + "input": "query", + "output": "response" + } + }, + { + "subset": "zh", + "columns": { + "instruction": "query", + "output": "response" + } + }], + "split": ["train", "val", "test"], + "tags": ["chat", "medical"] + }, + { + "ms_dataset_id": "swift/swift-sft-mixture", + "subsets": ["sharegpt", "firefly", "codefuse", "metamathqa"], + "tags": ["chat", "sft", "general", "🔥"], + "huge_dataset": true + }, + { + "ms_dataset_id": "ZhipuAI/LongWriter-6k", + "tags": ["long", "chat", "sft", "🔥"], + "hf_dataset_id": "zai-org/LongWriter-6k" + }, + { + "ms_dataset_id": "swift/longwriter-6k-filtered", + "tags": ["long", "chat", "sft", "🔥"] + }, + { + "ms_dataset_id": "AI-ModelScope/Magpie-Qwen2-Pro-300K-Filtered", + "tags": ["chat", "sft", "🔥"], + "hf_dataset_id": "Magpie-Align/Magpie-Qwen2-Pro-300K-Filtered" + }, + { + "ms_dataset_id": "AI-ModelScope/Magpie-Qwen2-Pro-200K-Chinese", + "tags": ["chat", "sft", "🔥", "zh"], + "hf_dataset_id": "Magpie-Align/Magpie-Qwen2-Pro-200K-Chinese" + }, + { + "ms_dataset_id": "AI-ModelScope/Magpie-Qwen2-Pro-200K-English", + "tags": ["chat", "sft", "🔥", "en"], + "hf_dataset_id": "Magpie-Align/Magpie-Qwen2-Pro-200K-English" + }, + { + "ms_dataset_id": "PowerInfer/QWQ-LONGCOT-500K", + "tags": ["chat", "sft", "🔥", "cot"], + "hf_dataset_id": "PowerInfer/QWQ-LONGCOT-500K" + }, + { + "ms_dataset_id": "PowerInfer/LONGCOT-Refine-500K", + "tags": ["chat", "sft", "🔥", "cot"], + "hf_dataset_id": "PowerInfer/LONGCOT-Refine-500K" + }, + { + "ms_dataset_id": "ServiceNow-AI/R1-Distill-SFT", + "hf_dataset_id": "ServiceNow-AI/R1-Distill-SFT", + "tags": ["chat", "sft", "cot", "r1"], + "subsets": [{ + "subset": "v0", + "columns": { + "problem": "query", + "reannotated_assistant_content": "response" + } + }, + { + "subset": "v1", + "columns": { + "messages": "_", + "reannotated_messages": "messages" + } + }] + }, + { + "ms_dataset_id": "bespokelabs/Bespoke-Stratos-17k", + "hf_dataset_id": "bespokelabs/Bespoke-Stratos-17k", + "tags": ["chat", "sft", "cot", "r1"] + }, + { + "ms_dataset_id": "open-thoughts/OpenThoughts-114k", + "hf_dataset_id": "open-thoughts/OpenThoughts-114k", + "tags": ["chat", "sft", "cot", "r1"] + }, + { + "ms_dataset_id": "HumanLLMs/Human-Like-DPO-Dataset", + "hf_dataset_id": "HumanLLMs/Human-Like-DPO-Dataset", + "columns": { + "prompt": "query", + "chosen": "response", + "rejected": "rejected_response" + }, + "tags": ["rlhf", "dpo"] + }, + { + "ms_dataset_id": "AI-ModelScope/MATH-lighteval", + "hf_dataset_id": "DigitalLearningGmbH/MATH-lighteval", + "columns": { + "problem": "query" + }, + "tags": ["grpo", "math"] + }, + { + "ms_dataset_id": "liucong/Chinese-DeepSeek-R1-Distill-data-110k-SFT", + "hf_dataset_id": "Congliu/Chinese-DeepSeek-R1-Distill-data-110k-SFT", + "tags": ["chat", "sft", "cot", "r1", "🔥"] + }, + { + "ms_dataset_id": "AI-MO/NuminaMath-CoT", + "hf_dataset_id": "AI-MO/NuminaMath-CoT", + "tags": ["grpo", "math"] + }, + { + "ms_dataset_id": "AI-MO/NuminaMath-TIR", + "hf_dataset_id": "AI-MO/NuminaMath-TIR", + "tags": ["grpo", "math", "🔥"] + }, + { + "ms_dataset_id": "AI-MO/NuminaMath-1.5", + "hf_dataset_id": "AI-MO/NuminaMath-1.5", + "tags": ["grpo", "math"] + }, + { + "ms_dataset_id": "FreedomIntelligence/medical-o1-reasoning-SFT", + "hf_dataset_id": "FreedomIntelligence/medical-o1-reasoning-SFT", + "subsets": ["en", "zh"], + "tags": ["medical", "o1", "🔥"] + }, + { + "ms_dataset_id": "lmms-lab/multimodal-open-r1-8k-verified", + "hf_dataset_id": "lmms-lab/multimodal-open-r1-8k-verified", + "tags": ["grpo", "vision", "🔥"] + }, + { + "ms_dataset_id": "open-r1/verifiable-coding-problems-python-10k", + "hf_dataset_id": "open-r1/verifiable-coding-problems-python-10k", + "tags": ["grpo", "code"] + }, + { + "ms_dataset_id": "open-r1/verifiable-coding-problems-python", + "hf_dataset_id": "open-r1/verifiable-coding-problems-python", + "columns": { + "problem_statement": "query" + }, + "tags": ["grpo", "code"] + }, + { + "ms_dataset_id": "open-r1/verifiable-coding-problems-python-10k_decontaminated", + "hf_dataset_id": "open-r1/verifiable-coding-problems-python-10k_decontaminated", + "tags": ["grpo", "code"] + }, + { + "ms_dataset_id": "open-r1/verifiable-coding-problems-python_decontaminated", + "hf_dataset_id": "open-r1/verifiable-coding-problems-python_decontaminated", + "columns": { + "problem_statement": "query" + }, + "tags": ["grpo", "code"] + }, + { + "ms_dataset_id": "iic/DocQA-RL-1.6K", + "hf_dataset_id": "Tongyi-Zhiwen/DocQA-RL-1.6K", + "columns": { + "prompt": "messages" + }, + "tags": ["docqa", "rl", "long-sequence"] + }, + { + "ms_dataset_id": "swift/Chinese-Qwen3-235B-2507-Distill-data-110k-SFT", + "tags": ["🔥", "distill", "sft"] + }, + { + "ms_dataset_id": "swift/Chinese-Qwen3-235B-Thinking-2507-Distill-data-110k-SFT", + "tags": ["🔥", "distill", "sft", "cot", "r1", "thinking"] + } +] diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2a74354eb1c136d5a1cbc31052be3ff4f646b335 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from . import llm, mllm diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/llm.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/llm.py new file mode 100644 index 0000000000000000000000000000000000000000..a144964f2213d8af7c8f7e79bf4e01222db2a36e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/llm.py @@ -0,0 +1,959 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import ast +import json +import numpy as np +import re +from functools import partial +from typing import Any, Dict, List, Optional, Tuple, Union + +from swift.template import split_str_parts_by +from ..preprocessor import (AlpacaPreprocessor, ClsGenerationPreprocessor, ClsPreprocessor, MessagesPreprocessor, + ResponsePreprocessor, RowPreprocessor, TextGenerationPreprocessor) +from ..register import DatasetMeta, SubsetDataset, register_dataset + + +class AlpacaZhPreprocessor(AlpacaPreprocessor): + + @classmethod + def concat_inst_input(cls, instruction, input_): + if input_ and input_.startswith('输入：'): + input_ = input_[3:] + return super().concat_inst_input(instruction, input_) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/alpaca-gpt4-data-zh', + hf_dataset_id='llm-wizard/alpaca-gpt4-data-zh', + preprocess_func=AlpacaZhPreprocessor(), + tags=['chat', 'general', '🔥'], + )) + + +class LongAlpacaPreprocessor(AlpacaPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + response = row['response'] + prefix_prompt = 'Answer: ' + if response and response.startswith(prefix_prompt): + response = response[len(prefix_prompt):].strip() + row['output'] = response + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/LongAlpaca-12k', + hf_dataset_id='Yukang/LongAlpaca-12k', + preprocess_func=LongAlpacaPreprocessor(), + tags=['long-sequence', 'QA'], + )) + + +class RuozhibaPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + title = row['title'] if row.get('title', None) is not None else row['content'] + abs = row['abs'] if 'abs' in row else None + if abs and abs != title: + title = title + '，' + abs + + pattern = r'\d+[\.,\s,\、](.+)' + match = re.search(pattern, title) + if match: + title = match.group(1) + if title: + return {'messages': [{'role': 'assistant', 'content': title}]} + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/ruozhiba', + subsets=['post-annual', 'title-good', 'title-norm'], + preprocess_func=RuozhibaPreprocessor(), + tags=['pretrain', '🔥'])) + + +class MathTrnPreprocessor(ResponsePreprocessor): + + def preprocess(self, row): + query = row['query'] + output = row['response'] + row = { + 'query': query, + 'response': output, + } + return super().preprocess(row) + + +register_dataset( + DatasetMeta(ms_dataset_id='AI-ModelScope/math-trn-format', preprocess_func=MathTrnPreprocessor(), tags=['math'])) + + +def _repair_ms_bench(messages: str) -> Optional[List[Dict[str, str]]]: + if isinstance(messages, str): + messages = ast.literal_eval(messages) + default_system = 'You are a helpful assistant.' + messages: List[Dict[str, str]] + if messages[0]['from'] == 'system' and messages[0]['value'] == default_system: + messages.pop(0) + # skip MOSS + for c in messages: + value = c['value'].lower() + if 'moss' in value or 'human:' in value or 'assistant:' in value or 'user:' in value: + return + return messages + + +register_dataset( + DatasetMeta( + ms_dataset_id='iic/ms_bench', + preprocess_func=MessagesPreprocessor(repair_messages=_repair_ms_bench), + tags=['chat', 'general', 'multi-round', '🔥'])) + + +def _repair_agent_messages(messages: List[Dict[str, str]], use_mini: bool) -> Optional[List[Dict[str, str]]]: + if use_mini: + pattern = r'\d\. {"plugin_name": "(.+?)"' + if messages[0]['from'] != 'system': + return + system = messages[0]['value'] + find_list = re.findall(pattern, system) + if len(set(find_list)) <= 1: + return + return messages + + +register_dataset( + DatasetMeta( + ms_dataset_id='damo/MSAgent-Bench', + subsets=[ + SubsetDataset( + preprocess_func=MessagesPreprocessor(repair_messages=partial(_repair_agent_messages, use_mini=False))), + SubsetDataset( + name='mini', + preprocess_func=MessagesPreprocessor(repair_messages=partial(_repair_agent_messages, use_mini=True)), + is_weak_subset=True) + ], + split=['train', 'validation'], + tags=['chat', 'agent', 'multi-round'])) + +advertise_gen_prompt = """Task: Generating advertisements based on keywords. +Keywords: {{QUERY}} +Advertisements:""" + +register_dataset( + DatasetMeta( + ms_dataset_id='lvjianjin/AdvertiseGen', + hf_dataset_id='shibing624/AdvertiseGen', + preprocess_func=TextGenerationPreprocessor( + prompt=advertise_gen_prompt, columns={ + 'content': 'query', + 'summary': 'response' + }), + tags=['text-generation', '🔥'], + split=['train', 'validation'], + )) + + +class FireflyPreprocessor(ResponsePreprocessor): + _firefly_kind_list = { + 'ProseGeneration', 'MRC', 'JinYongGeneration', 'TextCorrection', 'ClassicalChinese', 'BELLE', 'StoryGeneration', + 'Couplet', 'Cot', 'Dictionary', 'Translation', 'Program', 'SentimentAnalyze', 'OpenQA', 'AncientPoem', + 'TextMatching', 'NLI', 'Summary', 'KeywordRecognition', 'ProductDesc', 'LyricGeneration', 'Composition', + 'MusicComment', 'NER' + } + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if row['kind'] not in FireflyPreprocessor._firefly_kind_list: + return + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/firefly-train-1.1M', + hf_dataset_id='YeungNLP/firefly-train-1.1M', + preprocess_func=FireflyPreprocessor(), + tags=['chat', 'general'], + )) + +register_dataset( + DatasetMeta( + ms_dataset_id='modelscope/clue', + hf_dataset_id='clue', + subsets=['cmnli'], + preprocess_func=ClsGenerationPreprocessor(['neutral', 'entailment', 'contradiction'], + task='Natural Language Inference', + is_pair_seq=True), + tags=['text-generation', 'classification'], + split=['train', 'validation'], + )) + +register_dataset( + DatasetMeta( + ms_dataset_id='DAMO_NLP/jd', + subsets=[ + SubsetDataset( + 'default', + 'default', + preprocess_func=ClsGenerationPreprocessor(['negative', 'positive'], + task='Sentiment Classification', + is_pair_seq=False)), + SubsetDataset( + 'cls', + 'default', + preprocess_func=ClsPreprocessor(columns={'sentence': 'query'}), + ), + ], + tags=['text-generation', 'classification', '🔥'], + split=['train', 'validation'], + )) + + +class SyntheticText2SqlPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + sql_prompt = row['sql_prompt'] + sql_context = row['sql_context'] + sql = row['sql'] + sql_explanation = row['sql_explanation'] + query = f'Sql Table information:\n{sql_context}\n{sql_prompt}' + response = f'Let\'s think step by step:\n{sql_explanation}\nSo the final sql is:\n{sql}' + return super().preprocess({'query': query, 'response': response}) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/synthetic_text_to_sql', + hf_dataset_id='gretelai/synthetic_text_to_sql', + preprocess_func=SyntheticText2SqlPreprocessor(), + tags=['nl2sql', 'en'])) + + +def _repair_toolbench(conversations: List[Dict[str, str]]) -> List[Dict[str, str]]: + assert len(conversations) == 2 + if conversations[1]['from'] in {'caller', 'conclusion'}: + conversations[1]['from'] = 'assistant' + return conversations + + +register_dataset( + DatasetMeta( + ms_dataset_id='shenweizhou/alpha-umi-toolbench-processed-v2', + subsets=['backbone', 'caller', 'planner', 'summarizer'], + preprocess_func=MessagesPreprocessor(repair_messages=_repair_toolbench), + tags=['chat', 'agent', '🔥'], + huge_dataset=True)) + + +class BlossomMathPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + output, answer = row['output'], row['answer'] + return super().preprocess({'query': row['query'], 'response': f'{output}\n\nAnswer: {answer}'}) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/blossom-math-v2', + hf_dataset_id='Azure99/blossom-math-v2', + preprocess_func=BlossomMathPreprocessor(), + tags=['chat', 'math', '🔥'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/sql-create-context', + hf_dataset_id='b-mc2/sql-create-context', + preprocess_func=AlpacaPreprocessor(columns={ + 'question': 'instruction', + 'context': 'input', + 'answer': 'output' + }), + tags=['chat', 'sql', '🔥'])) + + +class TigerBotLawPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + prompt = """{type} +{title} +""" + cur_prompt = prompt.format(type=row['type'], title=row['title']) + for i in range(1, 4): + chapter = row[f'chapter{i}'] + if chapter is not None: + cur_prompt += f'{chapter}' + cur_prompt += f'{row["response"]}' + return super().preprocess({'response': cur_prompt}) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/tigerbot-law-plugin', + hf_dataset_id='TigerResearch/tigerbot-law-plugin', + preprocess_func=TigerBotLawPreprocessor(), + tags=['text-generation', 'law', 'pretrained'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='codefuse-ai/CodeExercise-Python-27k', + preprocess_func=MessagesPreprocessor(columns={'chat_rounds': 'messages'}), + tags=['chat', 'coding', '🔥'])) + + +class LeetcodePythonPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + code_with_problem = row['code_with_problem'] + idx = code_with_problem.find('```python') + problem = code_with_problem[:idx] + if problem.startswith('# '): + problem = problem[2:] + code = code_with_problem[idx:].strip() + explanation = row['explanation_only'] + return super().preprocess({'query': problem, 'response': f'{code}\n\n{explanation}'}) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/leetcode-solutions-python', + preprocess_func=LeetcodePythonPreprocessor(), + tags=['chat', 'coding', '🔥'])) + + +class StsbPreprocessor(RowPreprocessor): + + def __init__(self, sim_threshold: Optional[float] = None): + self.sim_threshold = sim_threshold + super().__init__() + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row = { + 'messages': [{ + 'role': 'user', + 'content': row['sentence1'] + }], + 'positive_messages': [[{ + 'role': 'user', + 'content': row['sentence2'] + }]], + 'label': row['score'], + } + if self.sim_threshold is None or float(row['label']) >= self.sim_threshold: + return row + else: + return None + + +class StsbGeneratePreprocessor(ResponsePreprocessor): + prompt = """Task: Based on the given two sentences, provide a similarity score between 0.0 and 1.0. +Sentence 1: {text1} +Sentence 2: {text2} +Similarity score: """ + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + return super().preprocess({ + 'query': self.prompt.format(text1=row['sentence1'], text2=row['sentence2']), + 'response': f"{row['score']:.1f}" + }) + + +class StsbRegressionPreprocessor(StsbGeneratePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + return super(StsbGeneratePreprocessor, self).preprocess({ + 'query': + self.prompt.format(text1=row['sentence1'], text2=row['sentence2']), + 'label': + row['score'] + }) + + +register_dataset( + DatasetMeta( + ms_dataset_id='sentence-transformers/stsb', + hf_dataset_id='sentence-transformers/stsb', + subsets=[ + SubsetDataset('default', preprocess_func=StsbPreprocessor()), # embedding + SubsetDataset('positive', preprocess_func=StsbPreprocessor(sim_threshold=0.75)), # infonce + SubsetDataset('generate', preprocess_func=StsbGeneratePreprocessor()), + SubsetDataset('reg', preprocess_func=StsbRegressionPreprocessor()), + ], + tags=['similarity', '🔥'])) + + +class MTEBRerankPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + query = row['query'] + positives = row['positive'] if isinstance(row['positive'], list) else [row['positive']] + negatives = row['negative'] if isinstance(row['negative'], list) else [row['negative']] + + messages = [{'role': 'user', 'content': query}] + positive_messages = [[{'role': 'assistant', 'content': positive}] for positive in positives] + negative_messages = [[{'role': 'assistant', 'content': negative}] for negative in negatives] + + return {'messages': messages, 'positive_messages': positive_messages, 'negative_messages': negative_messages} + + +register_dataset( + DatasetMeta( + ms_dataset_id='MTEB/scidocs-reranking', + hf_dataset_id='mteb/scidocs-reranking', + split=['validation', 'test'], + preprocess_func=MTEBRerankPreprocessor(), + tags=['rerank', '🔥'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='MTEB/stackoverflowdupquestions-reranking', + hf_dataset_id='mteb/stackoverflowdupquestions-reranking', + split=['train', 'test'], + preprocess_func=MTEBRerankPreprocessor(), + tags=['rerank', '🔥'])) + + +def _repair_conversations_agent_instruct(s: str) -> List[Dict[str, Any]]: + s = s.replace('}\n {', '},\n {') + if isinstance(s, str): + s = ast.literal_eval(s) + return s + + +register_dataset( + DatasetMeta( + ms_dataset_id='huangjintao/AgentInstruct_copy', + subsets=['alfworld', 'db', 'kg', 'mind2web', 'os', 'webshop'], + preprocess_func=MessagesPreprocessor(repair_messages=_repair_conversations_agent_instruct), + tags=['chat', 'agent', 'multi-round'])) + + +class MultiRoleAgentPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + conv = row['conversations'] + res_prompt = '\n\n【注意事项】\n1. 这是聊天室，不要发送私信给任何人\n2. 仅代表你个人说话,不要扮演其他人，只根据对话历史进行回复\n3. 长话短说，不要说太多话，不要超过50字 ' + history_prompt = '\n\n【chat history】' + conv_prompt = '\n {name}:{content}' + query, response = '', conv[-1]['value'] + system = conv[0]['value'] if conv[0]['from'] == 'system' else '' + if conv[0]['from'] == 'user': + query = conv[0]['value'] + elif 'next_speakers:' not in system: + if '【注意事项】' not in system and system: + system += res_prompt + system += history_prompt + system += ''.join([conv_prompt.format(name=c['from'], content=c['value']) for c in conv[1:-1]]) + + if not query or not response: + return + + return { + 'messages': [{ + 'role': 'system', + 'content': system + }, { + 'role': 'user', + 'content': query + }, { + 'role': 'assistant', + 'content': response + }], + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='iic/MSAgent-MultiRole', + preprocess_func=MultiRoleAgentPreprocessor(), + tags=['chat', 'agent', 'multi-round', 'role-play', 'multi-agent'])) + +register_dataset(DatasetMeta(ms_dataset_id='swift/ToolBench', tags=['chat', 'agent', 'multi-round'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='tastelikefeet/competition_math', + subsets=[ + SubsetDataset( + name='default', + subset='default', + split=['train', 'test'], + ), + ], + tags=['qa', 'math'])) + +register_dataset(DatasetMeta(ms_dataset_id='modelscope/gsm8k', subsets=['main'], split=['train'], tags=['qa', 'math'])) + +register_dataset( + DatasetMeta(ms_dataset_id='modelscope/MathR', subsets=['default', 'clean'], split=['train'], tags=['qa', 'math'])) + +register_dataset( + DatasetMeta(ms_dataset_id='modelscope/MathR-32B-Distill', subsets=['data'], split=['train'], tags=['qa', 'math'])) + + +class CoundownTaskPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + numbers = row['nums'] + target = row.pop('response', None) + query = (f'Using the numbers {numbers}, create an equation that equals {target}.\n' + 'You can use basic arithmetic operations (+, -, *, /) and each number can only be used once.\n' + 'Show your work in tags. And return the final equation and answer ' + 'in tags, for example (1 + 2) / 3 * 4 = 4 .') + row.update({'target': target, 'query': query}) + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='zouxuhong/Countdown-Tasks-3to4', + subsets=['default'], + preprocess_func=CoundownTaskPreprocessor(), + tags=['math'])) + + +class HC3Preprocessor(ResponsePreprocessor): + prompt = """Classification Task: Are the following responses from a human or from ChatGPT? +Question: {question} +Answer: {answer} +Category: Human, ChatGPT +Output:""" + + def preprocess(self, row): + rows = [] + for response in ['Human', 'ChatGPT']: + query = self.prompt.format( + question=row['query'], answer=self.random_state.choice(row[f'{response.lower()}_answers'])) + rows.append(super().preprocess({'query': query, 'response': response})) + return rows + + +class HC3ClsPreprocessor(HC3Preprocessor): + + def preprocess(self, row): + rows = [] + for i, response in enumerate(['Human', 'ChatGPT']): + query = self.prompt.format( + question=row['query'], answer=self.random_state.choice(row[f'{response.lower()}_answers'])) + rows.append(ResponsePreprocessor.preprocess(self, {'query': query, 'label': i})) + return rows + + +hc3_subset_names = ['baike', 'open_qa', 'nlpcc_dbqa', 'finance', 'medicine', 'law', 'psychology'] +hc3_subsets: List[SubsetDataset] = [] +for hc3_subset_name in hc3_subset_names: + hc3_subsets.append( + SubsetDataset( + name=hc3_subset_name, + subset=hc3_subset_name, + preprocess_func=HC3Preprocessor(), + )) + hc3_subsets.append( + SubsetDataset( + name=f'{hc3_subset_name}_cls', + subset=hc3_subset_name, + preprocess_func=HC3ClsPreprocessor(), + )) + +register_dataset( + DatasetMeta( + ms_dataset_id='simpleai/HC3-Chinese', + hf_dataset_id='Hello-SimpleAI/HC3-Chinese', + subsets=hc3_subsets, + tags=['text-generation', 'classification', '🔥'])) + +hc3_subset_names = ['finance', 'medicine'] +hc3_subsets: List[SubsetDataset] = [] +for hc3_subset_name in hc3_subset_names: + hc3_subsets.append( + SubsetDataset( + name=hc3_subset_name, + subset=hc3_subset_name, + preprocess_func=HC3Preprocessor(), + )) + hc3_subsets.append( + SubsetDataset( + name=f'{hc3_subset_name}_cls', + subset=hc3_subset_name, + preprocess_func=HC3ClsPreprocessor(), + )) + +register_dataset( + DatasetMeta( + ms_dataset_id='simpleai/HC3', + hf_dataset_id='Hello-SimpleAI/HC3', + subsets=hc3_subsets, + preprocess_func=HC3Preprocessor(), + tags=['text-generation', 'classification', '🔥'])) + + +class DureaderPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + prompt = """Task: Question Generation +Context: {context} +Answer: {answer} +Question:""" + answer, context = row['text1'].split('[SEP]') + return { + 'messages': [{ + 'role': 'user', + 'content': prompt.format(context=context, answer=answer) + }, { + 'role': 'assistant', + 'content': row['text2'] + }] + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='modelscope/DuReader_robust-QG', + preprocess_func=DureaderPreprocessor(), + split=['train', 'validation', 'test'], + tags=['text-generation', '🔥'])) + + +class HHRLHFPreprocessor(RowPreprocessor): + + @staticmethod + def _to_messages(data): + messages = [] + for query, response in zip(data[::2], data[1::2]): + messages.append({'role': 'user', 'content': query}) + messages.append({'role': 'assistant', 'content': response}) + return messages + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + chosen = row['chosen'].strip() + rejected = row['rejected'].strip() + parts_chosen = [s.strip() for s in re.split('\n\nHuman:|\n\nAssistant:|\n\nHum:', chosen)] + parts_rejected = [s.strip() for s in re.split('\n\nHuman:|\n\nAssistant:|\n\nHum:', rejected)] + if parts_chosen[0].startswith('Human:'): + assert parts_rejected[0].startswith('Human:') + parts_chosen[0] = parts_chosen[0][6:].strip() + parts_rejected[0] = parts_rejected[0][6:].strip() + row['messages'] = self._to_messages(parts_chosen) + row['rejected_messages'] = self._to_messages(parts_rejected) + return row + + +# TODO meta file broken +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/hh-rlhf', + subsets=['helpful-base', 'helpful-online', 'helpful-rejection-sampled'], + preprocess_func=HHRLHFPreprocessor(), + split=['train', 'test'], + tags=['rlhf', 'dpo'], + huge_dataset=True)) + + +class XlamFunctionCallingPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + messages = [{'role': 'user', 'content': row['query']}] + response = row['answers'] + response = json.loads(response) + messages += [{'role': 'tool_call', 'content': json.dumps(content)} for content in response] + return {'messages': messages, 'tools': row['tools']} + + +class XlamFunctionCallingGRPOPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + query = row['query'] + answers = row['response'] + if isinstance(answers, str): + answers = json.loads(answers) + answer = np.random.choice(answers) + name = answer['name'] + args = json.dumps(answer['arguments']) + response = f'Action: {name}\nAction Input: {args}' + row = {'query': query, 'response': response, 'solution': response, 'tools': row['tools']} + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='LLM-Research/xlam-function-calling-60k', + hf_dataset_id='Salesforce/xlam-function-calling-60k', + subsets=[ + SubsetDataset('default', 'dataset', preprocess_func=XlamFunctionCallingPreprocessor()), + SubsetDataset('grpo', 'dataset', preprocess_func=XlamFunctionCallingGRPOPreprocessor()) + ], + tags=['agent', 'grpo', '🔥'])) + + +class HHRLHFCNPreprocessor(MessagesPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['messages'].append(row.pop('chosen')) + row['rejected_response'] = row['rejected']['text'] + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/hh_rlhf_cn', + subsets=['hh_rlhf', 'harmless_base_cn', 'harmless_base_en', 'helpful_base_cn', 'helpful_base_en'], + preprocess_func=HHRLHFCNPreprocessor(columns={'context': 'messages'}, content_key='text'), + split=['train', 'test'], + tags=['rlhf', 'dpo', '🔥'])) + + +def repair_conversations(s: Union[str, Any]) -> Any: + if isinstance(s, str): + s = s.replace('}\n {', '},{') + s = s.replace('}\n{', '},{') + s = s.replace('}{', '},{') + s = s.replace('}\n {', '},{') + return ast.literal_eval(s) + return s + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/lmsys-chat-1m', + hf_dataset_id='lmsys/lmsys-chat-1m', + preprocess_func=MessagesPreprocessor(repair_messages=repair_conversations), + tags=['chat', 'em'])) + + +class EmojiPreprocessr(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + # Remove dirty characters + row['query'] = row['query'].replace('️', '') + row['response'] = row['response'].replace('️', '') + row['rejected_response'] = row['rejected_response'].replace('️', '') + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='hjh0119/shareAI-Llama3-DPO-zh-en-emoji', + hf_dataset_id='shareAI/DPO-zh-en-emoji', + preprocess_func=EmojiPreprocessr(columns={ + 'answer_zh': 'response', + 'answer_en': 'rejected_response' + }), + tags=['rlhf', 'dpo'])) + +register_dataset( + DatasetMeta(ms_dataset_id='AI-ModelScope/ultrafeedback-binarized-preferences-cleaned-kto', tags=['rlhf', 'kto'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='OmniData/Zhihu-KOL-More-Than-100-Upvotes', + hf_dataset_id='bzb2023/Zhihu-KOL-More-Than-100-Upvotes', + tags=['zhihu', 'qa'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='OmniData/Zhihu-KOL', + hf_dataset_id='wangrui6/Zhihu-KOL', + huge_dataset=True, + tags=['zhihu', 'qa'], + )) + + +class GuanacoPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + instruction = row['instruction'] + input = row['input'] + output = row['output'] + history = [] + if instruction: + parts = split_str_parts_by( + instruction, ['User:', 'User：', 'Assistant：', 'Assistant:', 'Asssistent:', 'Assistent:', 'Assistenz:']) + for idx, part in enumerate(parts): + if idx % 2 == 0: + if 'user' not in part['key'].lower(): + return + history.append([part['content'], None]) + else: + if 'assist' not in part['key'].lower() and 'asssist' not in part['key'].lower(): + return + history[-1][-1] = part['content'] + if input.startswith('User:'): + input = input[len('User:'):].strip() + if any([not h[0] or not h[1] for h in history]): + return + + messages = [] + for h in history: + messages.append({'role': 'user', 'content': h[0]}) + messages.append({'role': 'assistant', 'content': h[1]}) + messages.append({'role': 'user', 'content': input}) + messages.append({'role': 'assistant', 'content': output}) + return { + 'messages': messages, + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/GuanacoDataset', + hf_dataset_id='JosephusCheung/GuanacoDataset', + preprocess_func=GuanacoPreprocessor(), + tags=['chat', 'zh'])) + + +class FunctionCallChatmlPreprocessor(MessagesPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + res = super().preprocess(row) + + if res['function_description']: + res['tools'] = res['function_description'].split('\n\n') + messages = res['messages'] + if messages[0]['role'] == 'system': + messages.pop(0) + return res + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/function-calling-chatml', + hf_dataset_id='Locutusque/function-calling-chatml', + preprocess_func=FunctionCallChatmlPreprocessor(), + tags=['agent', 'en', 'sft', '🔥'])) + + +class Dolly15kPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + instruction = row['instruction'] + context = row['context'] + response = row['response'] + query = '' + if context: + query = 'Here gives some useful information:\n' + query += context + query += '\n' + query += instruction + return { + 'messages': [{ + 'role': 'user', + 'content': query + }, { + 'role': 'assistant', + 'content': response + }], + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/databricks-dolly-15k', + hf_dataset_id='databricks/databricks-dolly-15k', + preprocess_func=Dolly15kPreprocessor(), + tags=['multi-task', 'en', 'quality'])) + + +class OrpoDPOMix40kPreprocessor(MessagesPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if row['source'] == 'toxic-dpo-v0.2': + return + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/orpo-dpo-mix-40k', + hf_dataset_id='mlabonne/orpo-dpo-mix-40k', + preprocess_func=OrpoDPOMix40kPreprocessor(columns={ + 'chosen': 'messages', + 'rejected': 'rejected_messages' + }), + tags=['dpo', 'orpo', 'en', 'quality'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/sharegpt', + subsets=['common-zh', 'unknow-zh', 'common-en'], + tags=['chat', 'general', 'multi-round'])) + + +class SelfCognitionPreprocessor(ResponsePreprocessor): + + def __init__(self, *args, query_suffix: str = '', response_prefix: str = '', **kwargs): + self.query_suffix = query_suffix + self.response_prefix = response_prefix + self.name: Optional[Tuple[str, str]] = None + self.author: Optional[Tuple[str, str]] = None + super().__init__(*args, **kwargs) + + def set_name_author(self, name, author): + self.name = name + self.author = author + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + for key in ['name', 'author']: + val = getattr(self, key) + if val is None: + continue + val = val[0] if row['tag'] == 'zh' else val[1] + if val is None: + continue + placeholder = '{{' + key.upper() + '}}' + row['query'] = row['query'].replace(placeholder, val) + row['response'] = row['response'].replace(placeholder, val) + + row['query'] = row['query'] + self.query_suffix + row['response'] = self.response_prefix + row['response'] + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/self-cognition', + hf_dataset_id='modelscope/self-cognition', + subsets=[ + SubsetDataset(preprocess_func=SelfCognitionPreprocessor()), + SubsetDataset( + 'qwen3', + preprocess_func=SelfCognitionPreprocessor( + query_suffix=' /no_think', response_prefix='\n\n\n\n')), + SubsetDataset( + 'empty_think', preprocess_func=SelfCognitionPreprocessor(response_prefix='\n\n\n\n')), + ], + dataset_name='self-cognition', + tags=['chat', 'self-cognition', '🔥'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='open-r1/DAPO-Math-17k-Processed', + hf_dataset_id='open-r1/DAPO-Math-17k-Processed', + subsets=['all'], + tags=['math', 'rlvr'])) + + +class SudokuPreprocessor(ResponsePreprocessor): + prompt = ('Solve the following 9x9 Sudoku puzzle. ' + "Empty cells are marked with '0'. " + 'Provide the completed grid as your answer.\n\n' + 'Puzzle:\n{puzzle}') + + @staticmethod + def _format_grid(s: str) -> str: + return '\n'.join(s[i:i + 9] for i in range(0, len(s), 9)) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + puzzle = row['query'].replace('.', '0') + response = row['response'] + puzzle = self._format_grid(puzzle) + response = self._format_grid(response) + return super().preprocess({'query': self.prompt.format(puzzle=puzzle), 'response': response}) + + +register_dataset( + DatasetMeta( + ms_dataset_id='sapientinc/sudoku-extreme-1k', + hf_dataset_id='sapientinc/sudoku-extreme-1k', + preprocess_func=SudokuPreprocessor(), + )) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/mllm.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/mllm.py new file mode 100644 index 0000000000000000000000000000000000000000..540d90264d32d145394f7d920b97b33f2fb895a8 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset/mllm.py @@ -0,0 +1,1391 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import ast +import numpy as np +import os +import zipfile +from datasets import Dataset as HfDataset +from datasets import IterableDataset as HfIterableDataset +from tqdm import tqdm +from typing import Any, Dict, List, Optional + +from swift.utils import download_file, get_hf_endpoint, use_hf_hub +from ..media import MediaResource +from ..preprocessor import GroundingMixin, MessagesPreprocessor, ResponsePreprocessor, RowPreprocessor +from ..register import DatasetMeta, SubsetDataset, register_dataset + + +class ShareGPT4oPreprocessor(MessagesPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + row = super().preprocess(row) + image = row['images'] + if not image: + return + image = os.path.join(self.prefix_path, image) + if not os.path.exists(image): + return + row['images'] = [image] + return row + + def prepare_dataset(self, dataset): + if not use_hf_hub(): + url = ('https://www.modelscope.cn/api/v1/datasets/AI-ModelScope/ShareGPT-4o/repo?' + 'Revision=master&FilePath=images.zip') + else: + url = f'{get_hf_endpoint()}/datasets/OpenGVLab/ShareGPT-4o/blob/main/images.zip' + local_dir = MediaResource.download(url, 'sharegpt_4o_images') + self.prefix_path = os.path.join(local_dir, 'mnt', 'petrelfs', 'wangwenhai', 'workspace_cef', '4o', 'image') + return super().prepare_dataset(dataset) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/ShareGPT-4o', + hf_dataset_id='OpenGVLab/ShareGPT-4o', + preprocess_func=ShareGPT4oPreprocessor(), + subsets=['image_caption'], + split=['images'], + tags=['vqa', 'multi-modal'], + )) + + +class GPT4vDataset(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['query'] = 'What is the caption of this image?' + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/gpt4v-dataset', + hf_dataset_id='laion/gpt4v-dataset', + preprocess_func=GPT4vDataset(columns={ + 'link': 'images', + 'caption': 'response' + }), + split=['train'], + tags=['en', 'caption', 'multi-modal', 'quality'], + huge_dataset=True, + )) + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/RLAIF-V-Dataset', + hf_dataset_id='openbmb/RLAIF-V-Dataset', + preprocess_func=ResponsePreprocessor(columns={ + 'question': 'query', + 'chosen': 'response', + 'rejected': 'rejected_response' + }), + tags=['rlhf', 'dpo', 'multi-modal', 'en'], + )) + + +class GarbagePreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['query'] = 'Task: Classify household waste.' + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='tany0699/garbage265', + preprocess_func=GarbagePreprocessor(columns={ + 'category': 'label', + 'image:FILE': 'images' + }), + tags=['cls', '🔥', 'multi-modal'], + )) + + +class SA1BPairedCaptionPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + prompt = ['图片中展示了什么', '讲述一下图片中内容', '告诉我里面有什么', '图片内容是啥'] + response = row['global_caption'] + query = np.random.choice(prompt) + return { + 'messages': [{ + 'role': 'user', + 'content': query, + }, { + 'role': 'assistant', + 'content': response, + }] + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='Tongyi-DataEngine/SA1B-Paired-Captions-Images', + preprocess_func=SA1BPairedCaptionPreprocessor(columns={ + 'opensource_url': 'images', + }), + tags=['zh', 'multi-modal', 'vqa'], + )) + + +class SA1BDenseCaptionPreprocessor(RowPreprocessor): + column_mapping = { + 'url': 'images', + } + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + prompt = ['图片中展示了什么', '讲述一下图片中内容', '告诉我里面有什么', '图片内容是啥'] + response = ast.literal_eval(row['cap_seg']) + response = response.get('global_caption') + query = np.random.choice(prompt) + return { + 'messages': [{ + 'role': 'user', + 'content': query, + }, { + 'role': 'assistant', + 'content': response, + }] + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='Tongyi-DataEngine/SA1B-Dense-Caption', + preprocess_func=SA1BDenseCaptionPreprocessor(columns={ + 'url': 'images', + }), + tags=['zh', 'multi-modal', 'vqa'], + huge_dataset=True, + )) + + +class COCO2014Preprocess(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + caption = row['caption'] + if '&&' in caption: + caption = caption.split('&&')[0] + row['query'] = 'please describe the image.' + row['response'] = caption + + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='modelscope/coco_2014_caption', + preprocess_func=COCO2014Preprocess(), + subsets=[ + SubsetDataset('train', 'coco_2014_caption', ['train']), + SubsetDataset('validation', 'coco_2014_caption', ['validation']), + ], + tags=['chat', 'multi-modal', 'vision', '🔥'], + )) + + +class MantisPreprocessor(MessagesPreprocessor): + + def __init__(self, *, subset: str, columns: Optional[Dict[str, str]] = None) -> None: + self.subset = subset + super().__init__(columns=columns) + + def prepare_dataset(self, dataset: HfDataset) -> HfDataset: + if not use_hf_hub(): + url = (f'https://www.modelscope.cn/api/v1/datasets/swift/Mantis-Instruct/repo?Revision=' + f'master&FilePath={self.subset}/train_images.zip') # noqa + else: + url = (f'{get_hf_endpoint()}/datasets/TIGER-Lab/Mantis-Instruct/' + f'resolve/main/{self.subset}/train_images.zip') + self.local_dir = MediaResource.download(url, f'mantis_{self.subset}') + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + images = [os.path.join(self.local_dir, p['path']) for p in row['images']] + if not all([os.path.exists(d) for d in images]): + images = [] + + if not images: + return + row['images'] = images + return super().preprocess(row) + + +mantis_subsets_name = [ + 'birds-to-words', 'chartqa', 'coinstruct', 'contrastive_caption', 'docvqa', 'dreamsim', 'dvqa', 'iconqa', + 'imagecode', 'llava_665k_multi', 'lrv_multi', 'multi_vqa', 'nextqa', 'nlvr2', 'spot-the-diff', 'star', + 'visual_story_telling' +] + +_mantis_subsets = [] +for subset in mantis_subsets_name: + _subset = SubsetDataset(subset=subset, split=['train'], preprocess_func=MantisPreprocessor(subset=subset)) + _mantis_subsets.append(_subset) + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/Mantis-Instruct', + subsets=_mantis_subsets, + tags=['chat', 'multi-modal', 'vision'], + )) + + +class LLaVADataPreprocessor(MessagesPreprocessor): + + def prepare_dataset(self, dataset): + self.all_folders = {} + for media_type in ['coco', 'gqa', 'ocr_vqa', 'textvqa', 'VG_100K', 'VG_100K_2']: + self.all_folders[media_type] = MediaResource.download(media_type) + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if not row['images']: + return + row = super().preprocess(row) + images = [p['path'] for p in row['images']] + new_images = [] + for image in images: + if 'coco/' in image: + image = os.path.join(self.all_folders['coco'], image.replace('coco/', '')) + elif 'gqa/' in image: + image = os.path.join(self.all_folders['gqa'], image.replace('gqa/', '')) + elif 'ocr_vqa/' in image: + image = os.path.join(self.all_folders['ocr_vqa'], image) + elif 'textvqa/' in image: + image = os.path.join(self.all_folders['textvqa'], image.replace('textvqa/', '')) + elif 'VG_100K/' in image: + image = os.path.join(self.all_folders['VG_100K'], image.replace('vg/', '')) + elif 'VG_100K_2/' in image: + image = os.path.join(self.all_folders['VG_100K_2'], image.replace('vg/', '')) + new_images.append(image) + if all(os.path.exists(image) for image in new_images): + row['images'] = new_images + else: + return {'images': None} + return row + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/llava-data', + hf_dataset_id='TIGER-Lab/llava-data', + subsets=['llava_instruct'], + preprocess_func=LLaVADataPreprocessor(), + tags=['sft', 'multi-modal', 'quality'], + )) + + +class PixelProsePreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + caption_prompt = [ + 'Give the description of this image.', 'Describe this picture', 'What is the proper title of this image?' + ] + vlm_caption = row['vlm_caption'] + if vlm_caption.startswith('This image displays:'): + vlm_caption = vlm_caption[len('This image displays:'):].strip() + return { + 'messages': [{ + 'role': 'user', + 'content': np.random.choice(caption_prompt) + }, { + 'role': 'assistant', + 'content': vlm_caption + }], + 'images': + row['url'] + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/pixelprose', + hf_dataset_id='tomg-group-umd/pixelprose', + preprocess_func=PixelProsePreprocessor(), + split=['train', 'cc12m', 'commonpool', 'redcaps'], + tags=['caption', 'multi-modal', 'vision'], + huge_dataset=True, + )) + + +class AIShell1Preprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['query'] = '语音转文本' + row['response'] = row['Text:LABEL'].replace(' ', '') + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='speech_asr/speech_asr_aishell1_trainsets', + subsets=[ + SubsetDataset('train', split=['train']), + SubsetDataset('validation', split=['validation']), + SubsetDataset('test', split=['test']), + ], + preprocess_func=AIShell1Preprocessor(columns={'Audio:FILE': 'audios'}), + tags=['chat', 'multi-modal', 'audio'], + )) + + +class EmoSchemaPreprocessor(ResponsePreprocessor): + + def prepare_dataset(self, dataset: HfDataset) -> HfDataset: + for i in range(1, 6): + if not use_hf_hub(): + url = f'https://modelscope.cn/datasets/AI-ModelScope/egoschema/resolve/master/videos_chunked_0{i}.zip' + else: + url = f'{get_hf_endpoint()}/datasets/lmms-lab/egoschema/resolve/main/videos_chunked_0{i}.zip' + local_dir = MediaResource.download(url, 'egoschema') + + self.local_dir = os.path.join(local_dir, 'videos') + self.mp4_set = [file[:-4] for file in os.listdir(self.local_dir) if file.endswith('mp4')] + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if row['video_idx'] not in self.mp4_set: + return None + transfer_to_option = { + '0': 'A', + '1': 'B', + '2': 'C', + '3': 'D', + '4': 'E', + } + row = { + 'query': row['query'] + '\n' + '\n'.join(row['option']), + 'response': transfer_to_option[row['response']], + 'videos': [os.path.join(self.local_dir, f"{row['video_idx']}.mp4")], + } + return super().preprocess(row) + + +class EmoSchemaClsPreprocessor(EmoSchemaPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if row['video_idx'] not in self.mp4_set: + return None + row = { + 'query': row['query'] + '\n' + '\n'.join(row['option']), + 'label': int(row['response']), + 'videos': [os.path.join(self.local_dir, f"{row['video_idx']}.mp4")], + } + return ResponsePreprocessor.preprocess(self, row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/egoschema', + hf_dataset_id='lmms-lab/egoschema', + subsets=[ + SubsetDataset('default', 'Subset', preprocess_func=EmoSchemaPreprocessor()), + SubsetDataset('cls', 'Subset', preprocess_func=EmoSchemaClsPreprocessor()) + ], + split=['test'], + tags=['chat', 'multi-modal', 'video'], + )) + + +def _generate_url_list(_url, _range): + lst = [] + for i in range(1, (_range + 1)): + lst.append(_url.replace('{}', str(i))) + return lst + + +class LLaVAVideo178KPreprocessor(MessagesPreprocessor): + + def __init__(self, *, subset: str, columns: Optional[Dict[str, str]] = None) -> None: + self.subset = subset + super().__init__(columns=columns) + + url_prefix = 'https://www.modelscope.cn/datasets/lmms-lab/LLaVA-Video-178K/resolve/master/' + if use_hf_hub(): + url_prefix = f'{get_hf_endpoint()}/datasets/lmms-lab/LLaVA-Video-178K/resolve/main/' + + video_resources = { + '0_30_s_academic_v0_1': + _generate_url_list( + url_prefix + '0_30_s_academic_v0_1/0_30_s_academic_v0_1_videos_{}.tar.gz', + 8, + ), + '0_30_s_youtube_v0_1': + _generate_url_list( + url_prefix + '0_30_s_youtube_v0_1/0_30_s_youtube_v0_1_videos_{}.tar.gz', + 19, + ), + '1_2_m_academic_v0_1': + _generate_url_list( + url_prefix + '1_2_m_academic_v0_1/1_2_m_academic_v0_1_videos_{}.tar.gz', + 14, + ), + '1_2_m_youtube_v0_1': + _generate_url_list( + url_prefix + '1_2_m_youtube_v0_1/1_2_m_youtube_v0_1_videos_{}.tar.gz', + 50, + ), + '2_3_m_academic_v0_1': + _generate_url_list( + url_prefix + '2_3_m_academic_v0_1/2_3_m_academic_v0_1_videos_{}.tar.gz', + 18, + ), + '2_3_m_youtube_v0_1': + _generate_url_list( + url_prefix + '2_3_m_youtube_v0_1/2_3_m_youtube_v0_1_videos_{}.tar.gz', + 98, + ), + '30_60_s_academic_v0_1': + _generate_url_list( + url_prefix + '30_60_s_academic_v0_1/30_60_s_academic_v0_1_videos_{}.tar.gz', + 10, + ), + '30_60_s_youtube_v0_1': + _generate_url_list( + url_prefix + '30_60_s_youtube_v0_1/30_60_s_youtube_v0_1_videos_{}.tar.gz', + 13, + ), + } + + def prepare_dataset(self, dataset: HfDataset) -> HfDataset: + urls = self.video_resources[self.subset] + self.local_dir = MediaResource.download(urls, f'llava_video_178k_{self.subset}', file_type='sharded') + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + file_path = os.path.join(self.local_dir, f"{row['videos']}") + if not os.path.exists(file_path): + return None + return super().preprocess({'messages': row['messages'], 'videos': file_path}) + + +llava_video_subsets = [] +for subset in [ + '0_30_s_academic_v0_1', + '0_30_s_youtube_v0_1', + '1_2_m_academic_v0_1', + '1_2_m_youtube_v0_1', + '2_3_m_academic_v0_1', + '2_3_m_youtube_v0_1', + '30_60_s_academic_v0_1', + '30_60_s_youtube_v0_1', +]: + subset = SubsetDataset( + subset=subset, + split=['caption', 'open_ended', 'multi_choice'], + preprocess_func=LLaVAVideo178KPreprocessor(subset=subset), + ) + llava_video_subsets.append(subset) + +register_dataset( + DatasetMeta( + hf_dataset_id='lmms-lab/LLaVA-Video-178K', subsets=llava_video_subsets, tags=['chat', 'multi-modal', 'video'])) + + +class MovieChat1KPreprocessor(ResponsePreprocessor): + + def prepare_dataset(self, dataset: HfDataset) -> HfDataset: + mp4_set = [f'{i}.mp4' for i in range(1, 10)] + \ + [f'{i}.mp4' for i in range(201, 240)] + \ + [f'AWA-{i}.mp4' for i in range(1, 10)] + \ + [f'AWB-{i}.mp4' for i in range(1, 16)] + \ + [f'AWC-{i}.mp4' for i in range(1, 11)] + \ + [f'AWD-{i}.mp4' for i in range(1, 8)] + \ + [f'AWE-{i}.mp4' for i in range(1, 7)] + \ + [f'AWG-{i}.mp4' for i in range(1, 12)] + \ + [f'AWH-{i}.mp4' for i in range(1, 8)] + \ + [f'BWA-{i}.mp4' for i in range(1, 7)] + \ + [f'BWB-{i}.mp4' for i in range(1, 7)] + \ + [f'BWD-{i}.mp4' for i in range(1, 6)] + \ + [f'BWE-{i}.mp4' for i in range(1, 6)] + \ + [f'BWG-{i}.mp4' for i in range(1, 6)] + \ + [f'BWH-{i}.mp4' for i in range(1, 6)] + \ + [f'TFS-{i}.mp4' for i in range(1, 13)] + \ + [f'UWA-{i}.mp4' for i in range(1, 5)] + ['UWA-6.mp4'] + for file in mp4_set: + if not use_hf_hub(): + url = f'https://modelscope.cn/datasets/AI-ModelScope/MovieChat-1K-test/resolve/master/videos/{file}' + else: + url = f'{get_hf_endpoint()}/datasets/Enxin/MovieChat-1K-test/resolve/main/videos/{file}' + self.local_dir = MediaResource.download(url, 'moviechat_1k_test', file_type='file') + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + file_path = os.path.join(self.local_dir, f"{row['info']['video_path']}") + if not os.path.exists(file_path): + return None + return super().preprocess({ + 'query': row['global'][0]['question'], + 'response': row['global'][0]['answer'], + 'videos': file_path, + }) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/MovieChat-1K-test', + hf_dataset_id='Enxin/MovieChat-1K-test', + preprocess_func=MovieChat1KPreprocessor(), + split=['train'], + tags=['chat', 'multi-modal', 'video'])) + + +class VideoChatGPTPreprocessor(ResponsePreprocessor): + + def prepare_dataset(self, dataset: HfDataset) -> HfDataset: + if not use_hf_hub(): + url = 'https://modelscope.cn/datasets/swift/VideoChatGPT/resolve/master/videos.zip' + else: + url = f'{get_hf_endpoint()}/datasets/lmms-lab/VideoChatGPT/resolve/main/videos.zip' + local_dir = MediaResource.download(url, 'video_chatgpt') + self.local_dir = os.path.join(local_dir, 'Test_Videos') + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + # only `.mp4` + mp4_set = [file[:-4] for file in os.listdir(self.local_dir) if file.endswith('mp4')] + if row['video_name'] not in mp4_set: + return + row['videos'] = os.path.join(self.local_dir, f"{row['video_name']}.mp4") + for key in ['query', 'question_1', 'question_2']: + query = row.get(key) + if query is None or query == 'None': + continue + row['query'] = query + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/VideoChatGPT', + hf_dataset_id='lmms-lab/VideoChatGPT', + subsets=['Generic', 'Temporal', 'Consistency'], + preprocess_func=VideoChatGPTPreprocessor(), + split=['test'], + tags=['chat', 'multi-modal', 'video', '🔥'], + )) + + +def preprocess_mind2web(dataset, **kwargs): + + def preprocess_row(row: Dict[str, Any]) -> Dict[str, Any]: + raw_html = row['cleaned_html'] + screenshot = row['screenshot'] + row['screenshot'] = MediaResource.safe_save(screenshot, row['action_uid'] + '.jpg', 'mind2web') + action = row['target_action_reprs'] + actions = action.split('->') + row['query'] = f'The snapshot of screen:\nThe html source code:{raw_html}\n' + action = actions[-1] + where = actions[0] if len(actions) > 1 else '' + what = '' + if ':' in action: + action, what = action[:action.find(':')], action[action.find(':') + 1:] + row['response'] = f'Action: {action.strip()}\nAction Input: {where.strip()}{"," + what.strip()}' + return row + + conversations = [] + tools = [{ + 'function': { + 'name': 'CLICK', + 'desc': 'Choose and click an element in the web page', + 'parameter': [{ + 'element': 'string, the element in the web page to click' + }] + } + }, { + 'function': { + 'name': + 'TYPE', + 'desc': + 'Input some text into a web element like or ', + 'parameter': [{ + 'element': 'string, the element in the web page to input to', + 'content': 'string, what content to input into the textbox element' + }] + } + }, { + 'function': { + 'name': + 'SELECT', + 'desc': + 'Select an element from a combobox', + 'parameter': [{ + 'element': 'string, the combobox or dropdown in the web page on which the select happens', + 'content': 'string, which choices to choose' + }] + } + }] + + def history_to_messages(history): + messages = [] + for h in history: + messages.append({'role': 'user', 'content': h[0]}) + messages.append({'role': 'assistant', 'content': h[1]}) + return messages + + if isinstance(dataset, HfIterableDataset): + + def generate_example(dataset): + history = [] + images = [] + for row in dataset: + target_action_index = row['target_action_index'] + row = preprocess_row(row) + query = row['query'] + if target_action_index == '0': + if history: + yield {'messages': history_to_messages(history), 'images': images, 'tools': tools} + images = [] + history = [] + query = query + '\n' + row['confirmed_task'] + history.append([query, row['response']]) + images.append(row['screenshot']) + + if history: + yield {'messages': history_to_messages(history), 'images': images, 'tools': tools} + + return HfIterableDataset.from_generator(generate_example, gen_kwargs={'dataset': dataset}) + + history = [] + images = [] + for row in tqdm(dataset): + target_action_index = row['target_action_index'] + row = preprocess_row(row) + query = row['query'] + if target_action_index == '0': + if history: + conversations.append({'messages': history_to_messages(history), 'images': images, 'tools': tools}) + images = [] + history = [] + query = query + '\n' + row['confirmed_task'] + history.append([query, row['response']]) + images.append(row['screenshot']) + + if history: + conversations.append({'messages': history_to_messages(history), 'images': images, 'tools': tools}) + + return HfDataset.from_list(conversations) + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/Multimodal-Mind2Web', + hf_dataset_id='osunlp/Multimodal-Mind2Web', + preprocess_func=preprocess_mind2web, + tags=['agent', 'multi-modal'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/M3IT', + subsets=[ + 'coco', 'vqa-v2', 'shapes', 'shapes-rephrased', 'coco-goi-rephrased', 'snli-ve', 'snli-ve-rephrased', + 'okvqa', 'a-okvqa', 'viquae', 'textcap', 'docvqa', 'science-qa', 'imagenet', 'imagenet-open-ended', + 'imagenet-rephrased', 'coco-goi', 'clevr', 'clevr-rephrased', 'nlvr', 'coco-itm', 'coco-itm-rephrased', + 'vsr', 'vsr-rephrased', 'mocheg', 'mocheg-rephrased', 'coco-text', 'fm-iqa', 'activitynet-qa', 'msrvtt', + 'ss', 'coco-cn', 'refcoco', 'refcoco-rephrased', 'multi30k', 'image-paragraph-captioning', 'visual-dialog', + 'visual-dialog-rephrased', 'iqa', 'vcr', 'visual-mrc', 'ivqa', 'msrvtt-qa', 'msvd-qa', 'gqa', 'text-vqa', + 'ocr-vqa', 'st-vqa', 'flickr8k-cn' + ], + preprocess_func=ResponsePreprocessor(columns={ + 'instruction': 'system', + 'inputs': 'query', + 'image_base64_str': 'images', + 'outputs': 'response' + }), + split=['train'], + huge_dataset=True, + tags=['chat', 'multi-modal', 'vision'])) + + +class ShareGPT4VPreprocessor(MessagesPreprocessor): + + def prepare_dataset(self, dataset): + split = ['ShareGPT4V', 'ShareGPT4V-PT'] if dataset.config_name is None else dataset.config_name + IMAGE_DATASET_REQUIREMENTS = { + 'ShareGPT4V': ['coco', 'sam', 'llava', 'wikiart', 'share_textvqa', 'web-celebrity', 'web-landmark'], + 'ShareGPT4V-PT': ['coco', 'sam', 'llava'] + } + + if isinstance(split, str): + split = [split] + self.all_folders = {} + for sp in split: + for media_type in IMAGE_DATASET_REQUIREMENTS[sp]: + self.all_folders[media_type] = MediaResource.download(media_type) + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + image = row['image'] + row.update(super().preprocess(row)) + if 'coco/' in image: + image = os.path.join(self.all_folders['coco'], image.replace('coco/', '')) + elif 'sam/' in image: + image = os.path.join(self.all_folders['sam'], image.replace('sam/images/', '')) + elif 'llava/' in image: + image = os.path.join(self.all_folders['llava'], image.replace('llava/llava_pretrain/images/', '')) + elif 'wikiart/' in image: + image = os.path.join(self.all_folders['wikiart'], image.replace('wikiart/images/', 'data/wikiart/images/')) + elif 'share_textvqa/' in image: + image = os.path.join(self.all_folders['share_textvqa'], + image.replace('share_textvqa/images/', 'data/share_textvqa/images/')) + elif 'web-celebrity/' in image: + image = os.path.join(self.all_folders['web-celebrity'], + image.replace('web-celebrity/images/', 'data/web-celebrity/images/')) + elif 'web-landmark/' in image: + image = os.path.join(self.all_folders['web-landmark'], + image.replace('web-landmark/images/', 'data/web-landmark/images/')) + if os.path.exists(image): + row['images'] = image + else: + return + return row + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/ShareGPT4V', + subsets=['ShareGPT4V', 'ShareGPT4V-PT'], + preprocess_func=ShareGPT4VPreprocessor(), + huge_dataset=True, + tags=['chat', 'multi-modal', 'vision'])) + + +class TextCapsPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + row['query'] = 'What is the caption of this image?' + if not os.path.exists(row['images']['path']): + return None + return super().preprocess(row) + + +class TextCapsEmbPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if not os.path.exists(row['images']['path']): + return None + return { + 'messages': [{ + 'role': 'user', + 'content': '' + }], + 'positive_messages': [[{ + 'role': 'user', + 'content': row['response'][0] + }]], + 'images': row['images'] + } + + +class TextCapsReRankPreprocessor(RowPreprocessor): + + def __init__(self, + *, + columns: Optional[Dict[str, str]] = None, + negatives_per_sample: Optional[int] = None, + **kwargs): + super().__init__(columns=columns, **kwargs) + self._responses_pool: Optional[List[str]] = None + # Keep default aligned with MAX_NEGATIVE_SAMPLES used in collator if not provided + self.negatives_per_sample: int = int(os.environ.get( + 'MAX_NEGATIVE_SAMPLES', 7)) if negatives_per_sample is None else negatives_per_sample + + def prepare_dataset(self, dataset): + # Build a pool of responses from the dataset for negative sampling + try: + # Access full column; works for map-style datasets + responses = dataset['response'] if 'response' in dataset.features else None + except Exception: + responses = None + + pool: List[str] = [] + if responses is not None: + for resp in responses: + if isinstance(resp, (list, tuple)): + for s in resp: + if isinstance(s, str): + pool.append(s) + elif isinstance(resp, str): + pool.append(resp) + self._responses_pool = pool if pool else None + return dataset + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if not os.path.exists(row['images']['path']): + return None + positive = row['response'][0] + negatives: List[str] = [] + if self._responses_pool: + idxs = self.random_state.permutation(len(self._responses_pool)) + for i in idxs: + s = self._responses_pool[i] + if s not in row['response']: + negatives.append(s) + if len(negatives) >= self.negatives_per_sample: + break + return { + 'messages': [{ + 'role': 'user', + 'content': '' + }], + 'positive_messages': [[{ + 'role': 'user', + 'content': positive + }]], + 'negative_messages': [[{ + 'role': 'user', + 'content': n + }] for n in negatives], + 'images': row['images'] + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/TextCaps', + hf_dataset_id='HuggingFaceM4/TextCaps', + subsets=[ + SubsetDataset( + name='default', + preprocess_func=TextCapsPreprocessor(columns={'reference_strs': 'response'}), + split=['train', 'validation'], + ), + SubsetDataset( + name='emb', + preprocess_func=TextCapsEmbPreprocessor(columns={'reference_strs': 'response'}), + split=['train', 'validation'], + ), + SubsetDataset( + name='rerank', + preprocess_func=TextCapsReRankPreprocessor(columns={'reference_strs': 'response'}), + split=['train', 'validation'], + ), + ], + huge_dataset=True, + tags=['multi-modal', 'en', 'caption', 'quality'])) + + +class RefCOCOPreprocessor(ResponsePreprocessor, GroundingMixin): + task_type = 'caption' + + def __init__(self, task_type, **kwargs): + self.task_type = task_type + super().__init__(**kwargs) + + def prepare_dataset(self, dataset): + self.cache_dir = MediaResource.download( + 'https://www.modelscope.cn/api/v1/datasets/we_dont_produce_water/' + 'coco_res/repo?Revision=master&FilePath=coco_2014.zip', 'coco2014') + return dataset + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + caption = row['captions'][0] + bbox = row['bbox'] + image_path = os.path.join(self.cache_dir, row['image_path'].replace('coco/train2014', 'train2014')) + if not os.path.exists(image_path): + return + + for i in range(len(bbox)): + bbox[i] = round(float(bbox[i])) + res = {} + + objects = { + 'ref': [caption], + 'bbox': [bbox], + } + res['query'], res['response'] = self.construct_grounding_prompt() + res['images'] = [image_path] + res['objects'] = objects + return super().preprocess(res) + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/refcoco', + hf_dataset_id='jxu124/refcoco', + subsets=[ + SubsetDataset( + name='caption', + preprocess_func=RefCOCOPreprocessor('caption'), + ), + SubsetDataset( + name='grounding', + preprocess_func=RefCOCOPreprocessor('grounding'), + ) + ], + split=['train', 'validation'], + tags=['multi-modal', 'en', 'grounding'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/refcocog', + hf_dataset_id='jxu124/refcocog', + subsets=[ + SubsetDataset( + name='caption', + preprocess_func=RefCOCOPreprocessor('caption'), + ), + SubsetDataset( + name='grounding', + preprocess_func=RefCOCOPreprocessor('grounding'), + ) + ], + split=['train', 'validation'], + tags=['multi-modal', 'en', 'grounding'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/lnqa', + hf_dataset_id='vikhyatk/lnqa', + preprocess_func=MessagesPreprocessor(user_role='question', assistant_role='answer'), + split=['train', 'validation'], + huge_dataset=True, + tags=['multi-modal', 'en', 'ocr-vqa', 'quality'])) + + +class LLaVAInstructPreprocessor(MessagesPreprocessor): + + def prepare_dataset(self, dataset): + self.all_folders = {} + for media_type in ['coco', 'gqa', 'ocr_vqa', 'textvqa', 'VG_100K', 'VG_100K_2']: + self.all_folders[media_type] = MediaResource.download(media_type) + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + image = row['images'] + if 'coco/' in image: + image = os.path.join(self.all_folders['coco'], image.replace('coco/', '')) + elif 'gqa/' in image: + image = os.path.join(self.all_folders['gqa'], image.replace('gqa/', '')) + elif 'ocr_vqa/' in image: + image = os.path.join(self.all_folders['ocr_vqa'], image) + elif 'textvqa/' in image: + image = os.path.join(self.all_folders['textvqa'], image.replace('textvqa/', '')) + elif 'VG_100K/' in image: + image = os.path.join(self.all_folders['VG_100K'], image.replace('vg/', '')) + elif 'VG_100K_2/' in image: + image = os.path.join(self.all_folders['VG_100K_2'], image.replace('vg/', '')) + if os.path.exists(image): + row['images'] = image + else: + return + + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/LLaVA-Instruct-150K', + ms_revision='d5db3806e395c60496630a206c336932e85a2d00', + preprocess_func=LLaVAInstructPreprocessor(), + split=['train'], + tags=['chat', 'multi-modal', 'vision'])) + + +class LLaVAPretrainPreprocessor(MessagesPreprocessor): + + def prepare_dataset(self, dataset): + if not use_hf_hub(): + url = ('https://www.modelscope.cn/api/v1/datasets/AI-ModelScope/LLaVA-Pretrain/repo?' + 'Revision=master&FilePath=images.zip') + else: + url = f'{get_hf_endpoint()}/datasets/liuhaotian/LLaVA-Pretrain/resolve/main/images.zip' + self.media_dir = MediaResource.download( + url, + # noqa + 'llava_pretrain') + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + row.update(super().preprocess(row)) + if row['image']: + file_path = os.path.join(self.media_dir, row['image']) + if os.path.exists(file_path): + return {'images': file_path} + else: + return + else: + return + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/LLaVA-Pretrain', + ms_revision='e3a3f0bfaad05e90e46745152a32bf944e0f4a63', + hf_dataset_id='liuhaotian/LLaVA-Pretrain', + preprocess_func=LLaVAPretrainPreprocessor(), + huge_dataset=True, + tags=['chat', 'multi-modal', 'quality'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/MideficsDataset', + hf_dataset_id='WinterSchool/MideficsDataset', + preprocess_func=MessagesPreprocessor(inner_key='data', user_role='question', assistant_role='answer'), + tags=['medical', 'en', 'vqa'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/OK-VQA_train', + hf_dataset_id='Multimodal-Fatima/OK-VQA_train', + preprocess_func=ResponsePreprocessor(), + tags=['multi-modal', 'en', 'vqa', 'quality'])) + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/A-OKVQA', + hf_dataset_id='HuggingFaceM4/A-OKVQA', + split=['train', 'validation'], + preprocess_func=ResponsePreprocessor(columns={'rationales': 'response'}), + tags=['multi-modal', 'en', 'vqa', 'quality'])) + + +class OcrvqaPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + idx = np.random.choice(range(len(row['questions']))) + query = row['questions'][idx] + response = row['answers'][idx] + return { + 'messages': [{ + 'role': 'user', + 'content': query + }, { + 'role': 'assistant', + 'content': response + }], + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/OCR-VQA', + hf_dataset_id='howard-hou/OCR-VQA', + split=['train', 'validation'], + preprocess_func=OcrvqaPreprocessor(), + tags=['multi-modal', 'en', 'ocr-vqa'])) + + +class ScienceQAPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + query = row['question'] + response = row['choices'][row['answer']] + solution = row['solution'] + response = f'{solution}\nSo the final answer is: {response}' + return {'messages': [{'role': 'user', 'content': query}, {'role': 'assistant', 'content': response}]} + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/ScienceQA', + hf_dataset_id='derek-thomas/ScienceQA', + split=['train', 'validation'], + preprocess_func=ScienceQAPreprocessor(), + tags=['multi-modal', 'science', 'vqa', 'quality'])) + + +class GritPreprocessor(RowPreprocessor, GroundingMixin): + + def __init__(self, task_type, **kwargs): + self.task_type = task_type + super().__init__(**kwargs) + + @staticmethod + def has_overlap(start_ends): + for i in range(1, len(start_ends)): + if start_ends[i][0] < start_ends[i - 1][1]: + return True + return False + + @staticmethod + def replace_intervals_with_tags(response, start_ends): + result = [] + last_end = 0 + for start, end in start_ends: + result.append(response[int(last_end):int(start)]) + result.append('') + last_end = end + result.append(response[int(last_end):]) + return ''.join(result) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + images = row['images'] + caption = row['caption'] + ref_exps = row['ref_exps'] + objects = {'ref': [], 'bbox': [], 'bbox_type': 'norm1'} + start_end_pairs = [] + for ref_exp in ref_exps: + start = ref_exp[0] + end = ref_exp[1] + # conf = ref_exp[6] TODO filter low confidence rows? + start_end_pairs.append(ref_exp[0:2]) + + object_part = caption[int(start):int(end)] + objects['ref'].append(object_part) + objects['bbox'].append(ref_exp[2:6]) + + start_end_pairs.sort(key=lambda x: (x[0], x[1])) + if self.has_overlap(start_end_pairs) or not ref_exps: + return + + if self.task_type in ('grounding', 'caption'): + query, response = self.construct_grounding_prompt() + else: + query = 'what is the proper caption of this image?' + response = caption + return { + 'messages': [{ + 'role': 'user', + 'content': query + }, { + 'role': 'assistant', + 'content': response + }], + 'images': images, + 'objects': objects + } + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/GRIT', + hf_dataset_id='zzliang/GRIT', + subsets=[ + SubsetDataset( + name='caption', + preprocess_func=GritPreprocessor('caption', columns={'url': 'images'}), + ), + SubsetDataset( + name='grounding', + preprocess_func=GritPreprocessor('grounding', columns={'url': 'images'}), + ), + SubsetDataset( + name='vqa', + preprocess_func=GritPreprocessor('vqa', columns={'url': 'images'}), + ) + ], + huge_dataset=True, + tags=['multi-modal', 'en', 'caption-grounding', 'vqa', 'quality'])) + + +class GQAPreprocessor(RowPreprocessor): + + def prepare_dataset(self, dataset): + self.local_cache = MediaResource.download('gqa') + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if os.path.join(self.local_cache, 'images', row['imageId'] + '.jpg'): + return { + 'messages': [{ + 'role': 'user', + 'content': row['question'] + }, { + 'role': 'assistant', + 'content': row['fullAnswer'] + }], + 'images': + os.path.join(self.local_cache, 'images', row['imageId'] + '.jpg'), + } + else: + return + + +register_dataset( + DatasetMeta( + hf_dataset_id='lmms-lab/GQA', + split=['train_all_instructions'], + preprocess_func=GQAPreprocessor(), + huge_dataset=True, + tags=['multi-modal', 'en', 'vqa', 'quality'])) + + +class CocoPreprocessor(ResponsePreprocessor): + category = [ + 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', + 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', + 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', + 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', + 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', + 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', + 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', + 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' + ] + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + row['query'] = 'Task: Object Detection' + objects = row['objects'] + objects['ref'] = [self.category[c] for c in objects['category']] + row['response'] = '\n'.join([''] * len(objects['ref'])) + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/coco', + hf_dataset_id='detection-datasets/coco', + preprocess_func=CocoPreprocessor(), + huge_dataset=True, + tags=['multi-modal', 'en', 'vqa', 'quality'])) + + +class LLaVAMixSFTPreprocessor(RowPreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + messages = row['messages'] + rounds = [] + for msg in messages: + role = msg['role'] + content = msg['content'] + text = '' + for index in content: + if index['type'] == 'text': + text += index['text'] + elif index['type'] == 'image': + text += '' + + rounds.append({'role': role, 'content': text}) + + return {'messages': rounds} + + +register_dataset( + DatasetMeta( + ms_dataset_id='swift/llava-instruct-mix-vsft', + hf_dataset_id='HuggingFaceH4/llava-instruct-mix-vsft', + split=['test'], + preprocess_func=LLaVAMixSFTPreprocessor(), + tags=['multi-modal', 'en', 'vqa', 'quality'])) + + +class LatexocrPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['query'] = 'Using LaTeX to perform OCR on the image.' + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/LaTeX_OCR', + hf_dataset_id='linxy/LaTeX_OCR', + subsets=['default', 'human_handwrite', 'human_handwrite_print', 'synthetic_handwrite', 'small'], + preprocess_func=LatexocrPreprocessor(), + split=['train', 'validation', 'test'], + tags=['chat', 'ocr', 'multi-modal', 'vision'], + )) + + +class CapchaImagesPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['query'] = 'recognize the content.' + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/captcha-images', + split=['train', 'validation'], + preprocess_func=CapchaImagesPreprocessor(columns={'solution': 'response'}), + tags=['chat', 'multi-modal', 'vision'])) + + +class ClevrPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + query = row.get('query', '') + query = (f'{query} Output the thinking process in and ' + 'final answer (number) in tags.') + row.update({'query': query}) + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='AI-ModelScope/clevr_cogen_a_train', + hf_dataset_id='leonardPKU/clevr_cogen_a_train', + preprocess_func=ClevrPreprocessor(), + tags=['qa', 'math', 'vision', 'grpo'])) + + +class Voc2007MultilabelPreprocessor(ResponsePreprocessor): + CLASS_NAME = ('aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', + 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor') + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['query'] = f'多标签分类，类别包括：{list(self.CLASS_NAME)}' + row['label'] = [i for i, x in enumerate(row['npy']) if x == 1] + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='clip-benchmark/wds_voc2007_multilabel', + hf_dataset_id='clip-benchmark/wds_voc2007_multilabel', + preprocess_func=Voc2007MultilabelPreprocessor(columns={'webp': 'images'}), + tags=['multilabel', 'multi-modal'], + )) + + +class Geometry3KPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['solution'] = row['response'] + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + hf_dataset_id='hiyouga/geometry3k', + subsets=[ + SubsetDataset('train', split=['train']), + SubsetDataset('validation', split=['validation']), + SubsetDataset('test', split=['test']), + ], + preprocess_func=Geometry3KPreprocessor(), + tags=['multi-modal', 'en', 'math'])) + + +class Qwen3TTSPreprocessor(ResponsePreprocessor): + """Preprocessor for Qwen3-TTS SFT datasets. + + Input format: + {'audio': 'path.wav', 'text': '...', 'ref_audio': 'ref.wav'} + Output format: + {'messages': [{'role': 'assistant', 'content': '...'}], + 'audios': ['path.wav'], 'ref_audios': ['ref.wav']} + """ + + data_dir: Optional[str] = None + + def prepare_dataset(self, dataset): + url = 'https://modelscope.cn/datasets/qsdong/Qwen3-1.7-TTS-SFT-Furina/resolve/master/Furina.zip' + file_path = download_file(url) + extract_dir = file_path.rsplit('.', 1)[0] + if not os.path.exists(extract_dir): + with zipfile.ZipFile(file_path, 'r') as zf: + zf.extractall(extract_dir) + self.data_dir = extract_dir + return super().prepare_dataset(dataset) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if self.data_dir: + audios = row.get('audios') + if isinstance(audios, str) and not os.path.isabs(audios): + row['audios'] = os.path.join(self.data_dir, audios) + ref_audios = row.get('ref_audios') + if isinstance(ref_audios, str) and not os.path.isabs(ref_audios): + row['ref_audios'] = os.path.join(self.data_dir, ref_audios) + ref_audios = row.get('ref_audios') + if isinstance(ref_audios, str): + row['ref_audios'] = [ref_audios] + return super().preprocess(row) + + +register_dataset( + DatasetMeta( + ms_dataset_id='qsdong/Qwen3-1.7-TTS-SFT-Furina', + preprocess_func=Qwen3TTSPreprocessor(columns={'ref_audio': 'ref_audios'}), + tags=['chat', 'multi-modal', 'audio', 'tts'], + )) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset_meta.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..c7d01e90d782f0f841bb2858d723aec7e82cec1d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset_meta.py @@ -0,0 +1,202 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import numpy as np +import os +import shutil +from abc import ABC, abstractmethod +from copy import deepcopy +from dataclasses import dataclass, field +from datasets import Dataset as HfDataset +from datasets import concatenate_datasets, interleave_datasets +from modelscope.hub.api import ModelScopeConfig +from modelscope.utils.config_ds import MS_CACHE_HOME +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union + +from swift.utils import download_ms_file, get_logger, get_seed, safe_ddp_context +from .preprocessor import DATASET_TYPE, AutoPreprocessor +from .utils import sample_dataset + +PreprocessFunc = Callable[..., DATASET_TYPE] +logger = get_logger() + +if TYPE_CHECKING: + from .dataset_syntax import DatasetSyntax + + +@dataclass +class SubsetDataset: + # `Name` is used for matching subsets of the dataset, and `subset` refers to the subset_name on the hub. + name: Optional[str] = None + # If set to None, then subset is set to subset_name. + subset: str = 'default' + + # Higher priority. If set to None, the attributes of the DatasetMeta will be used. + split: Optional[List[str]] = None + preprocess_func: Optional[PreprocessFunc] = None + + # If the dataset specifies "all," weak subsets will be skipped. + is_weak_subset: bool = False + + def __post_init__(self): + if self.name is None: + self.name = self.subset + + def set_default(self, dataset_meta: 'DatasetMeta') -> 'SubsetDataset': + subset_dataset = deepcopy(self) + for k in ['split', 'preprocess_func']: + v = getattr(subset_dataset, k) + if v is None: + setattr(subset_dataset, k, deepcopy(getattr(dataset_meta, k))) + return subset_dataset + + +class BaseDatasetLoader(ABC): + + @abstractmethod + def load( + self, + dataset_syntax: Optional['DatasetSyntax'] = None, + dataset_meta: Optional['DatasetMeta'] = None, + *, + use_hf: Optional[bool] = None, + ) -> HfDataset: + pass + + @staticmethod + def download_ms_dataset(ms_dataset_id: str, files: List[str], force_download: bool = False) -> str: + """Download dataset from repo manually + Args: + ms_dataset_id: The dataset id of ModelScope + files: Which files to download + force_download: Force download or not + Returns: + The dataset dir + """ + assert isinstance(files, list) + url = f'http://www.modelscope.cn/api/v1/datasets/{ms_dataset_id}/repo?Revision=master&FilePath={{fpath}}' + cache_dir = os.path.join(MS_CACHE_HOME, 'datasets', ms_dataset_id, 'master') + local_dir = os.path.join(cache_dir, 'raw') + tmp_dir = os.path.join(cache_dir, 'tmp') + os.makedirs(local_dir, exist_ok=True) + os.makedirs(tmp_dir, exist_ok=True) + cookies = ModelScopeConfig.get_cookies() + with TemporaryDirectory(dir=tmp_dir) as temp_dir: + for remote_fpath in files: + url = url.format(fpath=remote_fpath) + temp_fpath = os.path.join(temp_dir, remote_fpath) + local_fpath = os.path.join(local_dir, remote_fpath) + if not force_download and os.path.exists(local_fpath): + continue + download_ms_file(url, temp_fpath, cookies) + shutil.copy2(temp_fpath, local_fpath) + + return local_dir + + @staticmethod + def concat_datasets(datasets: List[HfDataset]) -> Optional[HfDataset]: + if len(datasets) == 0: + return + if len(datasets) == 1: + return datasets[0] + return concatenate_datasets(datasets) + + @staticmethod + def interleave_datasets(datasets, *args, **kwargs): + if len(datasets) == 0: + return + if len(datasets) == 1: + return datasets[0] + return interleave_datasets(datasets, *args, **kwargs) + + @staticmethod + def shuffle_dataset(dataset, seed: int, buffer_size: int = 1000): + if isinstance(dataset, HfDataset): + with safe_ddp_context(None, True): + return dataset.shuffle(seed=seed) + else: + return dataset.shuffle(seed=seed, buffer_size=buffer_size) + + @staticmethod + def post_process( + train_dataset: DATASET_TYPE, + *, + dataset_sample: Optional[int] = None, + split_dataset_ratio: float = 0., + streaming: bool = False, + shuffle: bool = True, + random_state: Optional[np.random.RandomState] = None, + ) -> Tuple[DATASET_TYPE, Optional[DATASET_TYPE]]: + """Split into train/val datasets and perform dataset sampling.""" + assert dataset_sample is None or dataset_sample > 0 + assert 0 <= split_dataset_ratio <= 1 + if streaming: + if dataset_sample is None: + if split_dataset_ratio == 0: + val_dataset = None + elif split_dataset_ratio == 1: + train_dataset, val_dataset = None, train_dataset + else: + raise ValueError('The IterableDataset does not support splitting the training set ' + 'and validation set when dataset_sample is None.') + else: + # not shuffle + train_dataset = train_dataset.take(dataset_sample) + val_sample = int(dataset_sample * split_dataset_ratio) + val_dataset = None if val_sample == 0 else train_dataset.take(val_sample) + if val_sample: + train_dataset = train_dataset.skip(val_sample) + else: + if dataset_sample is None: + dataset_sample = len(train_dataset) + if split_dataset_ratio == 0: + train_dataset = sample_dataset(train_dataset, dataset_sample, shuffle, random_state) + val_dataset = None + elif split_dataset_ratio == 1: + train_dataset, val_dataset = None, train_dataset + val_sample = dataset_sample + # Avoid duplication in the val_dataset. + assert val_sample <= len(val_dataset), f'val_sample: {val_sample}, len(val_dataset): {len(val_dataset)}' + val_dataset = sample_dataset(val_dataset, val_sample, shuffle, random_state) + else: + # Avoid duplication in the val_dataset. + train_len = min(len(train_dataset), dataset_sample) + val_sample = max(int(train_len * split_dataset_ratio), 1) + train_sample = dataset_sample - val_sample + assert train_sample > 0 + with safe_ddp_context(None, True): + train_dataset, val_dataset = train_dataset.train_test_split( + test_size=val_sample, shuffle=shuffle, seed=get_seed(random_state)).values() + train_dataset = sample_dataset(train_dataset, train_sample, shuffle, random_state) + return train_dataset, val_dataset + + +@dataclass +class DatasetMeta: + ms_dataset_id: Optional[str] = None + hf_dataset_id: Optional[str] = None + dataset_path: Optional[str] = None # or dataset_dir + dataset_name: Optional[str] = None + ms_revision: Optional[str] = None + hf_revision: Optional[str] = None + + subsets: List[Union[SubsetDataset, str]] = field(default_factory=lambda: ['default']) + # Applicable to all subsets. + split: List[str] = field(default_factory=lambda: ['train']) + # First perform column mapping, then proceed with the preprocess_func. + preprocess_func: PreprocessFunc = field(default_factory=lambda: AutoPreprocessor()) + loader: Optional[BaseDatasetLoader] = None + + tags: List[str] = field(default_factory=list) + help: Optional[str] = None + huge_dataset: bool = False + + def __post_init__(self): + from .loader import DatasetLoader + if self.loader is None: + self.loader = DatasetLoader + for i, subset in enumerate(self.subsets): + if isinstance(subset, str): + self.subsets[i] = SubsetDataset(subset=subset) + + +DATASET_MAPPING: Dict[Tuple[str, str, str], DatasetMeta] = {} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset_syntax.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset_syntax.py new file mode 100644 index 0000000000000000000000000000000000000000..950a1e9fca6dc3b339e1493e05fbaecaa1f3f71d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/dataset_syntax.py @@ -0,0 +1,127 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import platform +import re +from dataclasses import dataclass, field +from typing import Dict, List, Literal, Optional, Tuple + +from .dataset_meta import DATASET_MAPPING, DatasetMeta + +_dataset_meta_mapping = None + + +@dataclass +class DatasetSyntax: + dataset: str + subsets: List[str] = field(default_factory=list) + dataset_sample: Optional[int] = None + use_hf: Optional[bool] = None + + def __post_init__(self): + if os.path.isfile(self.dataset): + self.dataset_type = 'path' + else: # dataset_id or dataset_dir + self.dataset_type = 'repo' + + def get_raw(self): + subsets = '/'.join(self.subsets) + dataset_sample = '' if self.dataset_sample is None else f'#{self.dataset_sample}' + return f'{self.dataset}{subsets}{dataset_sample}' + + @staticmethod + def _safe_split(s: str, + sep: str, + use_0: bool, + split_mode: Literal['left', 'right'] = 'left') -> Tuple[Optional[str], Optional[str]]: + """ + use_0: When the length of the part is 1, is it considered as part0 or part1. + split_mode: use split or rsplit + """ + if s is None or len(s) == 0: + return None, None + if split_mode == 'left': + part = s.split(sep, 1) + else: + part = s.rsplit(sep, 1) + if len(part) == 1: + if use_0: + part = part[0], None + else: + part = None, part[0] + else: + assert len(part) == 2 + return part + + @classmethod + def parse(cls, dataset: str) -> 'DatasetSyntax': + """Parse the dataset from the command line""" + # hf/ms::dataset_id or dataset_path:subset1/subset2/subset3#dataset_sample + if os.path.exists(dataset): + use_hf = None + else: + use_hf, dataset = cls._safe_split(dataset, '::', False) + if isinstance(use_hf, str): + use_hf = use_hf.lower() + use_hf = {'hf': True, 'ms': False}.get(use_hf) + if os.path.exists(dataset): + other, dataset_sample = dataset, None + else: + other, dataset_sample = cls._safe_split(dataset, '#', True, 'right') + if os.path.exists(other): + dataset, subsets = other, None + else: + dataset, subsets = cls._safe_split(other, ':', True) + + if subsets is not None: + subsets = [subset.strip() for subset in subsets.split('/')] + if dataset_sample is not None: + dataset_sample = int(dataset_sample) + return cls(dataset.strip(), subsets or [], dataset_sample, use_hf) + + def get_dataset_meta(self, use_hf: bool): + dataset_meta_mapping = self._get_dataset_meta_mapping() + dataset_type = self.dataset_type + if dataset_type == 'path': + dataset_meta = dataset_meta_mapping.get((dataset_type, self.dataset)) + else: + dataset_type = 'repo' if os.path.isdir(self.dataset) else {True: 'hf', False: 'ms'}[use_hf] + dataset_meta = dataset_meta_mapping.get((dataset_type, self.dataset)) + return dataset_meta or self._get_matched_dataset_meta(dataset_meta_mapping) or DatasetMeta() + + @staticmethod + def _get_dataset_meta_mapping() -> Dict[Tuple[str, str], DatasetMeta]: + global _dataset_meta_mapping + if _dataset_meta_mapping is not None: + return _dataset_meta_mapping + _dataset_meta_mapping = {} + for dataset_meta in DATASET_MAPPING.values(): + if dataset_meta.dataset_path is not None: + dataset_type = 'repo' if os.path.isdir(dataset_meta.dataset_path) else 'path' + _dataset_meta_mapping[(dataset_type, dataset_meta.dataset_path)] = dataset_meta + if dataset_meta.ms_dataset_id is not None: + _dataset_meta_mapping[('ms', dataset_meta.ms_dataset_id)] = dataset_meta + if dataset_meta.hf_dataset_id is not None: + _dataset_meta_mapping[('hf', dataset_meta.hf_dataset_id)] = dataset_meta + return _dataset_meta_mapping + + @staticmethod + def get_dataset_name(dataset_id: str) -> str: + # compat hf hub + dataset_id = dataset_id.rstrip('/') + match_ = re.search('/datasets--.+?--(.+?)/snapshots/', dataset_id) + if match_ is not None: + return match_.group(1) + + dataset_name = dataset_id.rsplit('/', 1)[-1] + if platform.system().lower() == 'windows': + dataset_name = dataset_name.rsplit('\\', 1)[-1] + return dataset_name + + def _get_matched_dataset_meta(self, dataset_meta_mapping): + suffix_dataset_meta_mapping = {} + for dataset_name, dataset_meta in dataset_meta_mapping.items(): + dataset_name = self.get_dataset_name(dataset_name[1]) + suffix_dataset_meta_mapping[dataset_name] = dataset_meta + dataset_name = self.get_dataset_name(self.dataset) + dataset_meta = suffix_dataset_meta_mapping.get(dataset_name) + return dataset_meta diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/indexed_dataset.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/indexed_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..0c05683e95fc95f02498aabc122151c8d9921439 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/indexed_dataset.py @@ -0,0 +1,132 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import bisect +import mmap +import os +import pickle +import threading +from modelscope.hub.utils.utils import get_cache_dir +from queue import Queue +from torch.utils.data import Dataset +from typing import Any, List + + +class IndexedDatasetBuilder: + CHUNK_SIZE = 1e10 + + def __init__(self, dataset_name: str): + self.cache_dir = IndexedDataset.get_cache_dir(dataset_name) + self.n_shard = 1 + self.bin_path = os.path.join(self.cache_dir, IndexedDataset.BIN_FNAME.format(0)) + self.idx_path = os.path.join(self.cache_dir, IndexedDataset.IDX_FNAME) + if os.path.exists(self.bin_path): + os.remove(self.bin_path) + self.bin_file = open(self.bin_path, 'ab') + self.length_list = [] + self.idx_list = [0] + self.shard_offset = [0] + self._thread = None + self._queue = Queue(maxsize=1000) + + def _write_worker(self): + while True: + items = self._queue.get() + if items is None: + break + bin_buffer = [] + for item in items: + item_buffer = pickle.dumps(item) + bin_buffer.append(item_buffer) + self.idx_list.append(self.idx_list[-1] + len(item_buffer)) + self.length_list.append(item['length']) + self.bin_file.write(b''.join(bin_buffer)) + offset = self.idx_list[-1] - self.shard_offset[-1] + if offset >= self.CHUNK_SIZE: + self.bin_file.close() + self.bin_path = os.path.join(self.cache_dir, IndexedDataset.BIN_FNAME.format(self.n_shard)) + self.shard_offset.append(self.shard_offset[-1] + offset) + self.n_shard += 1 + if os.path.exists(self.bin_path): + os.remove(self.bin_path) + self.bin_file = open(self.bin_path, 'ab') + + def add_items(self, items: List[Any]) -> None: + if self._thread is None: + self._thread = threading.Thread(target=self._write_worker, daemon=True) + self._thread.start() + self._queue.put(items) + + def finalize(self): + if self._thread is not None: + self._queue.put(None) + self._thread.join() + self.bin_file.close() + idx_obj = { + 'idx': self.idx_list, + 'length': self.length_list, + 'n_shard': self.n_shard, + 'shard_offset': self.shard_offset, + } + with open(self.idx_path, 'wb') as f: + pickle.dump(idx_obj, f) + + +class BinReader: + + def __init__(self, bin_path: str): + self.bin_path = bin_path + self.file = open(bin_path, 'rb') + try: + self.mm = mmap.mmap(self.file.fileno(), 0, access=mmap.ACCESS_READ) + except ValueError: + # For example, self.file is an empty file. + self.mm = None + + def read_buffer(self, offset: int, size: int) -> bytes: + if offset < 0 or size < 0 or offset + size > len(self.mm): + raise ValueError('Invalid offset or size') + return self.mm[offset:offset + size] + + def __del__(self): + if self.mm is not None: + self.mm.close() + self.file.close() + + +class IndexedDataset(Dataset): + BIN_FNAME = 'data-{:05d}.bin' + IDX_FNAME = 'data.idx' + + @staticmethod + def get_cache_dir(dataset_name: str): + cache_dir = os.getenv('PACKING_CACHE') or os.path.join(get_cache_dir(), 'tmp') + cache_dir = os.path.join(cache_dir, dataset_name) + os.makedirs(cache_dir, exist_ok=True) + assert dataset_name is not None, f'dataset_name: {dataset_name}' + return cache_dir + + def __init__(self, dataset_name: str): + self.dataset_name = dataset_name + cache_dir = self.get_cache_dir(dataset_name) + self.idx_path = os.path.join(cache_dir, self.IDX_FNAME) + with open(self.idx_path, 'rb') as f: + idx_obj = pickle.load(f) + self.idx_list = idx_obj['idx'] + self.length_list = idx_obj['length'] + self.n_shard = idx_obj['n_shard'] + self.shard_offset = idx_obj['shard_offset'] + self.bin_readers = [] + for i in range(self.n_shard): + bin_path = os.path.join(cache_dir, self.BIN_FNAME.format(i)) + self.bin_readers.append(BinReader(bin_path)) + + def __getitem__(self, index: int): + if index < 0: + index = index % len(self) + idx, idx_next = self.idx_list[index], self.idx_list[index + 1] + num_shard = bisect.bisect_right(self.shard_offset, idx) + offset = self.shard_offset[num_shard - 1] + buffer = self.bin_readers[num_shard - 1].read_buffer(idx - offset, idx_next - idx) + return pickle.loads(buffer) + + def __len__(self): + return len(self.idx_list) - 1 diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/loader.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/loader.py new file mode 100644 index 0000000000000000000000000000000000000000..11f00d48d7f5840dc9d735c9b95044b50958d769 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/loader.py @@ -0,0 +1,384 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import numpy as np +import os +from contextlib import nullcontext +from datasets import Dataset as HfDataset +from datasets import IterableDataset as HfIterableDataset +from datasets import load_dataset as hf_load_dataset +from functools import partial +from modelscope.hub.utils.utils import get_cache_dir +from typing import Dict, List, Literal, Optional, Tuple, Union + +from swift.hub import get_hub +from swift.utils import get_logger, get_seed, safe_ddp_context, use_hf_hub +from .dataset_meta import DATASET_TYPE, BaseDatasetLoader +from .dataset_syntax import DatasetSyntax +from .preprocessor import RowPreprocessor +from .register import DATASET_MAPPING, DatasetMeta, SubsetDataset + +logger = get_logger() + + +class DatasetLoader(BaseDatasetLoader): + + def __init__( + self, + num_proc: int = 1, + load_from_cache_file: bool = True, + streaming: bool = False, + hub_token: Optional[str] = None, + strict: bool = False, + download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists', + columns: Optional[Dict[str, str]] = None, + remove_unused_columns: bool = True, + disable_auto_column_mapping: bool = False, + ): + self.num_proc = num_proc + self.load_from_cache_file = load_from_cache_file + self.streaming = streaming + self.hub_token = hub_token + self.strict = strict + self.download_mode = download_mode + self.columns = columns + self.remove_unused_columns = remove_unused_columns + self.disable_auto_column_mapping = disable_auto_column_mapping + + def _load_dataset_path( + self, + dataset_path: str, + dataset_meta: DatasetMeta, + ) -> HfDataset: + ext = os.path.splitext(dataset_path)[1].lstrip('.') + file_type = {'jsonl': 'json', 'txt': 'text'}.get(ext) or ext + kwargs = {'split': 'train', 'streaming': self.streaming, 'num_proc': self.num_proc} + if file_type == 'csv': + kwargs['na_filter'] = False + with safe_ddp_context(None, True): + kwargs['cache_dir'] = os.path.join(get_cache_dir(), 'datasets') + dataset = hf_load_dataset(file_type, data_files=dataset_path, **kwargs) + if self.columns: + dataset = RowPreprocessor.safe_rename_columns(dataset, self.columns) + dataset = dataset_meta.preprocess_func( + dataset, + num_proc=self.num_proc, + load_from_cache_file=self.load_from_cache_file, + strict=self.strict, + enable_auto_mapping=not self.disable_auto_column_mapping) + if self.remove_unused_columns: + dataset = RowPreprocessor.remove_useless_columns(dataset) + return dataset + + def _load_repo_dataset( + self, + dataset_id: str, + subset: SubsetDataset, + *, + use_hf: Optional[bool] = None, + revision: Optional[str] = None, + ) -> HfDataset: + datasets = [] + if os.path.isdir(dataset_id): + retry = 1 + load_context = nullcontext + use_hf = True + dataset_str = f'Use local folder, dataset_dir: {dataset_id}' + # The dataset downloaded from modelscope will have an additional dataset_infos.json file. + with safe_ddp_context('dataset_infos_rename'): + dataset_infos_path = os.path.join(dataset_id, 'dataset_infos.json') + if os.path.isfile(dataset_infos_path): + os.rename(dataset_infos_path, f'{dataset_infos_path}_bak') + elif dataset_id.startswith('/'): + raise ValueError(f'The local path does not exist, dataset_id: `{dataset_id}`. ' + f'os.path.exists(dataset_id): {os.path.exists(dataset_id)}') + else: + retry = 3 + load_context = partial(safe_ddp_context, hash_id=dataset_id, use_barrier=True) + dataset_str_f = 'Downloading the dataset from {hub}, dataset_id: {dataset_id}' + if use_hf: + dataset_str = dataset_str_f.format(hub='HuggingFace', dataset_id=dataset_id) + else: + dataset_str = dataset_str_f.format(hub='ModelScope', dataset_id=dataset_id) + logger.info(dataset_str) + hub = get_hub(use_hf) + for split in subset.split: + i = 1 + with load_context(): + while True: + try: + dataset = hub.load_dataset( + dataset_id, + subset.subset, + split, + streaming=self.streaming, + revision=revision, + download_mode=self.download_mode, + hub_token=self.hub_token, + num_proc=self.num_proc) + except Exception as e: + if i == retry: + raise + i += 1 + logger.error(f'Dataset {dataset_id} load failed: subset_name={subset.subset},' + f'split={split} with error: {e}') + else: + break + if hasattr(dataset, '_hf_ds'): + dataset = dataset._hf_ds + if self.streaming and isinstance(dataset, HfDataset): + dataset = dataset.to_iterable_dataset() + if self.columns: + dataset = RowPreprocessor.safe_rename_columns(dataset, self.columns) + dataset = subset.preprocess_func( + dataset, + num_proc=self.num_proc, + load_from_cache_file=self.load_from_cache_file, + strict=self.strict, + enable_auto_mapping=not self.disable_auto_column_mapping) + if self.remove_unused_columns: + dataset = RowPreprocessor.remove_useless_columns(dataset) + datasets.append(dataset) + return self.concat_datasets(datasets) + + @staticmethod + def _select_subsets(subsets: List[str], dataset_meta: DatasetMeta) -> List[SubsetDataset]: + subset_mapping = {subset.name: subset for subset in dataset_meta.subsets} + subset_names = list(subset_mapping.keys()) + if not subsets: + if len(subset_names) <= 1: + subsets = subset_names + elif 'default' in subset_names: + subsets = ['default'] + else: + raise ValueError(f'Please provide subsets. available subsets: {subset_names}') + elif len(subsets) == 1 and subsets[0] == 'all' and 'all' not in subset_names: + subsets = [subset_name for subset_name in subset_names if not subset_mapping[subset_name].is_weak_subset] + + subsets = [ + subset_mapping[subset_name] if subset_name in subset_mapping else SubsetDataset(subset=subset_name) + for subset_name in subsets + ] + return [subset.set_default(dataset_meta) for subset in subsets] + + def load( + self, + dataset_syntax: Optional[DatasetSyntax] = None, + dataset_meta: Optional[DatasetMeta] = None, + *, + use_hf: Optional[bool] = None, + ) -> HfDataset: + if dataset_syntax.dataset_type == 'path': + dataset = self._load_dataset_path( + dataset_syntax.dataset, + dataset_meta=dataset_meta, + ) + else: + subsets: List[SubsetDataset] = self._select_subsets(dataset_syntax.subsets, dataset_meta) + revision = dataset_meta.hf_revision if use_hf else dataset_meta.ms_revision + datasets = [] + for subset in subsets: + dataset = self._load_repo_dataset( + dataset_syntax.dataset, + subset, + use_hf=use_hf, + revision=revision, + ) + datasets.append(dataset) + dataset = self.concat_datasets(datasets) + return dataset + + +def init_self_cognition_preprocessor( + dataset_meta: Optional[DatasetMeta], + model_name: Optional[Union[Tuple[str, str], List[str]]] = None, + model_author: Optional[Union[Tuple[str, str], List[str]]] = None, +) -> None: + from .dataset.llm import SelfCognitionPreprocessor + if dataset_meta is None or model_name is None and model_author is None: + return + kwargs = {} + # zh, en + for key in ['name', 'author']: + val = locals()[f'model_{key}'] + if isinstance(val, str): + val = [val] + if val is not None and val[0] is not None and (len(val) == 1 or val[1] is None): + val = (val[0], val[0]) + kwargs[key] = val + + preprocess_funcs = [dataset_meta.preprocess_func] + preprocess_funcs += [subset.preprocess_func for subset in dataset_meta.subsets if isinstance(subset, SubsetDataset)] + for preprocess_func in preprocess_funcs: + if isinstance(preprocess_func, SelfCognitionPreprocessor): + preprocess_func.set_name_author(**kwargs) + logger.info_once(f"SelfCognitionPreprocessor has been successfully configured with name: {kwargs['name']}, " + f"author: {kwargs['author']}.") + + +def _inject_dataset_routing_tag(dataset: DATASET_TYPE, ds_name: str) -> DATASET_TYPE: + """Inject ``dataset`` column for multi-teacher routing (constant per source dataset).""" + if isinstance(dataset, HfIterableDataset): + return dataset.map(lambda example: {**example, 'dataset': ds_name}) + return dataset.add_column('dataset', [ds_name] * len(dataset)) + + +def load_dataset( + datasets: Union[List[str], str], + *, + split_dataset_ratio: float = 0., + seed: Union[int, np.random.RandomState, None] = 42, + num_proc: int = 1, + load_from_cache_file: bool = True, + shuffle: bool = False, + streaming: bool = False, + interleave_prob: Optional[List[float]] = None, + stopping_strategy: Literal['first_exhausted', 'all_exhausted'] = 'first_exhausted', + shuffle_buffer_size: int = 1000, + use_hf: Optional[bool] = None, + hub_token: Optional[str] = None, + strict: bool = False, + download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists', + columns: Optional[Dict[str, str]] = None, # columns_mapping + remove_unused_columns: bool = True, + disable_auto_column_mapping: bool = False, + # self-cognition + model_name: Optional[Union[Tuple[str, str], List[str]]] = None, # zh, en + model_author: Optional[Union[Tuple[str, str], List[str]]] = None, +) -> Tuple[DATASET_TYPE, Optional[DATASET_TYPE]]: + """Load and preprocess datasets. + + This function provides a unified interface to load datasets from various sources (HuggingFace, + ModelScope, or local paths), with support for splitting, shuffling, streaming, and interleaving + multiple datasets. It also handles self-cognition dataset preprocessing for model training. + + Args: + datasets: Single dataset name or list of dataset names to load. Can use special syntax + for advanced configurations (e.g., 'dataset_name#1000' for sampling). + split_dataset_ratio: Ratio for splitting dataset into train/validation sets. + Value between 0 and 1. If 0, no validation split is created. Default: 0. + seed: Random seed for reproducibility. Can be an integer or numpy RandomState object. + If None, results will be non-deterministic. Default: 42. + num_proc: Number of processes to use for dataset preprocessing. Set to None for + streaming mode. Default: 1. + load_from_cache_file: Whether to load preprocessed data from cache if available. + Default: True. + shuffle: Whether to shuffle the dataset(s) after loading. Default: False. + streaming: Enable streaming mode for large datasets that don't fit in memory. + When True, num_proc is automatically set to None. Default: False. + interleave_prob: Probability weights for interleaving multiple datasets. Must have + same length as datasets list. If None, datasets are concatenated instead. Default: None. + stopping_strategy: Strategy when interleaving datasets of different lengths: + - 'first_exhausted': Stop when shortest dataset is exhausted + - 'all_exhausted': Continue until all datasets are exhausted + Default: 'first_exhausted'. + shuffle_buffer_size: Buffer size for shuffling in streaming mode. Larger values + provide better randomization but use more memory. Default: 1000. + use_hf: Force using HuggingFace Hub (True) or ModelScope (False). If None, + it is controlled by the environment variable `USE_HF`, which defaults to '0'. + Default: None. + hub_token: Authentication token for accessing private datasets on the hub. Default: None. + strict: If True, raise exceptions when encountering malformed data rows. + If False, skip invalid rows with warnings. Default: False. + download_mode: How to handle existing cached datasets: + - 'reuse_dataset_if_exists': Use cached version if available + - 'force_redownload': Always download fresh copy + Default: 'reuse_dataset_if_exists'. + columns: Manual column name mapping for datasets. Dictionary mapping source column + names to target column names (e.g., {'text': 'content'}). Default: None. + remove_unused_columns: Whether to remove columns not used in preprocessing. + Helps reduce memory usage. Default: True. + disable_auto_column_mapping: By default, column names in the dataset are automatically + mapped. This parameter disables that behavior + (the `columns` parameter remains effective), defaulting to `False`. + model_name: Model name for self-cognition task preprocessing. Can be a tuple of + (Chinese_name, English_name) or list of names. Default: None. + model_author: Model author for self-cognition task preprocessing. Can be a tuple of + (Chinese_author, English_author) or list of authors. Default: None. + + Returns: + A tuple of (train_dataset, val_dataset): + - train_dataset: The training dataset + - val_dataset: The validation dataset if split_dataset_ratio > 0, otherwise None + + Examples: + >>> # Load single dataset + >>> train_ds, val_ds = load_dataset('AI-ModelScope/alpaca-gpt4-data-zh', split_dataset_ratio=0.1) + + >>> # Load multiple datasets + >>> train_ds, _ = load_dataset( + ... ['AI-ModelScope/alpaca-gpt4-data-zh#500', 'swift/self-cognition#500'], + ... model_name=('我的模型', 'MyModel'), + ... model_author=('作者', 'Author') + ... ) + """ + init_self_cognition_preprocessor(DATASET_MAPPING.get('self-cognition'), model_name, model_author) + if isinstance(datasets, str): + datasets = [datasets] + if not isinstance(seed, np.random.RandomState): + seed = np.random.RandomState(seed) + if streaming: + num_proc = None + train_datasets = [] + val_datasets = [] + + use_hf_default = use_hf + if use_hf_default is None: + use_hf_default = True if use_hf_hub() else False + for dataset in datasets: + dataset_syntax = DatasetSyntax.parse(dataset) + use_hf = dataset_syntax.use_hf or use_hf_default + # compat dataset_name + if dataset_syntax.dataset in DATASET_MAPPING: + dataset_meta = DATASET_MAPPING[dataset_syntax.dataset] + if dataset_syntax.use_hf is None and dataset_meta.dataset_path is not None: + dataset_syntax.dataset = dataset_meta.dataset_path + dataset_syntax.dataset_type = 'path' + else: + dataset_syntax.dataset = dataset_meta.hf_dataset_id if use_hf else dataset_meta.ms_dataset_id + else: + dataset_meta = dataset_syntax.get_dataset_meta(use_hf) + loader = dataset_meta.loader( + num_proc=num_proc, + load_from_cache_file=load_from_cache_file, + streaming=streaming, + hub_token=hub_token, + strict=strict, + download_mode=download_mode, + columns=columns, # columns_mapping + remove_unused_columns=remove_unused_columns, + disable_auto_column_mapping=disable_auto_column_mapping, + ) + train_dataset = loader.load(dataset_syntax, dataset_meta, use_hf=use_hf) + train_dataset, val_dataset = loader.post_process( + train_dataset, + dataset_sample=dataset_syntax.dataset_sample, + split_dataset_ratio=split_dataset_ratio, + streaming=streaming, + shuffle=shuffle, + random_state=seed, + ) + if train_dataset is not None: + # Inject dataset_syntax.dataset as routing tag for multi-teacher + ds_name = dataset_syntax.dataset + train_dataset = _inject_dataset_routing_tag(train_dataset, ds_name) + train_datasets.append(train_dataset) + if val_dataset is not None: + ds_name = dataset_syntax.dataset + val_dataset = _inject_dataset_routing_tag(val_dataset, ds_name) + val_datasets.append(val_dataset) + + if interleave_prob is None: + train_datasets = loader.concat_datasets(train_datasets) + val_datasets = loader.concat_datasets(val_datasets) + else: + train_datasets = loader.interleave_datasets( + train_datasets, interleave_prob, seed=get_seed(seed), stopping_strategy=stopping_strategy) + val_datasets = loader.interleave_datasets( + val_datasets, interleave_prob, seed=get_seed(seed), stopping_strategy=stopping_strategy) + + if shuffle: + if train_datasets: + train_datasets = loader.shuffle_dataset( + train_datasets, seed=get_seed(seed), buffer_size=shuffle_buffer_size) + if val_datasets: + val_datasets = loader.shuffle_dataset(val_datasets, seed=get_seed(seed), buffer_size=shuffle_buffer_size) + return train_datasets, val_datasets diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/media.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/media.py new file mode 100644 index 0000000000000000000000000000000000000000..40221ba26a24cb67472ebe28fdc22327eb67b196 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/media.py @@ -0,0 +1,128 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import aiohttp +import os +import shutil +from modelscope.hub.utils.utils import get_cache_dir +from typing import List, Literal, Optional, Union + +from swift.utils import get_logger, safe_ddp_context + +logger = get_logger() + + +class MediaResource: + """A class to manage the resource downloading.""" + + cache_dir = os.path.join(get_cache_dir(), 'media_resources') + lock_dir = os.path.join(get_cache_dir(), 'lockers') + + media_type_urls = { + 'llava', 'coco', 'sam', 'gqa', 'ocr_vqa', 'textvqa', 'VG_100K', 'VG_100K_2', 'share_textvqa', 'web-celebrity', + 'web-landmark', 'wikiart' + } + + URL_PREFIX = 'https://www.modelscope.cn/api/v1/datasets/hjh0119/sharegpt4v-images/repo?Revision=master&FilePath=' + + @staticmethod + def get_url(media_type): + is_ocr_vqa = (media_type == 'ocr_vqa') + extension = 'tar' if is_ocr_vqa else 'zip' + return f'{MediaResource.URL_PREFIX}{media_type}.{extension}' + + @staticmethod + def download(media_type_or_url: Union[str, List[str]], + local_alias: Optional[str] = None, + file_type: Literal['compressed', 'file', 'sharded'] = 'compressed'): + """Download and extract a resource from a http link. + + Args: + media_type_or_url: `str` or List or `str`, Either belongs to the `media_type_urls` + listed in the class field, or a remote url to download and extract. + Be aware that, this media type or url needs to contain a zip or tar file. + local_alias: `Options[str]`, The local alias name for the `media_type_or_url`. If the first arg is a + media_type listed in this class, local_alias can leave None. else please pass in a name for the url. + The local dir contains the extracted files will be: {cache_dir}/{local_alias} + file_type: The file type, if is a compressed file, un-compressed the file, + if is an original file, only download it, if is a sharded file, download all files and extract. + + Returns: + The local dir contains the extracted files. + """ + media_file = media_type_or_url if isinstance(media_type_or_url, str) else media_type_or_url[0] + with safe_ddp_context(hash_id=media_file): + return MediaResource._safe_download( + media_type=media_type_or_url, media_name=local_alias, file_type=file_type) + + @staticmethod + def move_directory_contents(src_dir, dst_dir): + if not os.path.exists(dst_dir): + os.makedirs(dst_dir) + + for dirpath, dirnames, filenames in os.walk(src_dir): + relative_path = os.path.relpath(dirpath, src_dir) + target_dir = os.path.join(dst_dir, relative_path) + + if not os.path.exists(target_dir): + os.makedirs(target_dir) + + for file in filenames: + src_file = os.path.join(dirpath, file) + dst_file = os.path.join(target_dir, file) + shutil.move(src_file, dst_file) + + @staticmethod + def _safe_download(media_type: Union[str, List[str]], + media_name: Optional[str] = None, + file_type: Literal['compressed', 'file', 'sharded'] = 'compressed'): + media_name = media_name or media_type + assert isinstance(media_name, str), f'{media_name} is not a str' + if isinstance(media_type, str) and media_type in MediaResource.media_type_urls: + media_type = MediaResource.get_url(media_type) + + from datasets.download.download_manager import DownloadConfig, DownloadManager + final_folder = os.path.join(MediaResource.cache_dir, media_name) + + if file_type == 'file': + filename = media_type.split('/')[-1] + final_path = os.path.join(final_folder, filename) + if os.path.exists(final_path): # if the download thing is a file but not folder, + return final_folder # check whether the file exists + if not os.path.exists(final_folder): + os.makedirs(final_folder) # and make sure final_folder exists to contain it + else: + if os.path.exists(final_folder): + return final_folder + + logger.info('# #################Resource downloading#################') + logger.info('Downloading necessary resources...') + logger.info(f'Resource package: {media_type}') + logger.info(f'Extracting to local dir: {final_folder}') + logger.info('If the downloading fails or lasts a long time, ' + 'you can manually download the resources and extracting to the local dir.') + logger.info('Now begin.') + download_config = DownloadConfig(cache_dir=MediaResource.cache_dir) + download_config.storage_options = {'client_kwargs': {'timeout': aiohttp.ClientTimeout(total=86400)}} + if file_type == 'file': + filename = media_type.split('/')[-1] + final_path = os.path.join(final_folder, filename) + local_dirs = DownloadManager(download_config=download_config).download(media_type) + shutil.move(str(local_dirs), final_path) + elif file_type == 'compressed': + local_dirs = DownloadManager(download_config=download_config).download_and_extract(media_type) + shutil.move(str(local_dirs), final_folder) + else: + for media_url in media_type: + local_dirs = DownloadManager(download_config=download_config).download_and_extract(media_url) + MediaResource.move_directory_contents(str(local_dirs), final_folder) + logger.info('# #################Resource downloading finished#################') + return final_folder + + @staticmethod + def safe_save(image, file_name, folder, format='JPEG'): + folder = os.path.join(MediaResource.cache_dir, folder) + os.makedirs(folder, exist_ok=True) + file = os.path.join(folder, file_name) + if os.path.exists(file): + return file + image.save(file, format=format) + return file diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/packing.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/packing.py new file mode 100644 index 0000000000000000000000000000000000000000..de0d434ffab7fc414f01aa4636a478862e5ac169 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/packing.py @@ -0,0 +1,234 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import math +import multiprocessing as mp +import torch.distributed as dist +from itertools import chain +from torch.utils.data import Dataset, IterableDataset +from tqdm import tqdm +from typing import Optional + +from swift.template import MaxLengthError +from swift.utils import get_logger, is_dist, is_master, split_list + +logger = get_logger() + + +def calculate_matched_group(sequences, packing_length: int, is_finished: bool = True, strategy: str = 'binpack'): + if len(sequences) == 0: + return [], [] + if strategy == 'sequential': + # Order-preserving greedy packing (next-fit): keep a single open pack and flush it + # when the next sample doesn't fit, so the global sample order and pack boundaries + # follow the input order (a sequential sampler). (Use packing_num_proc=1 for + # a single global ordering.) + packs, cur, cur_len = [], [], 0 + for item in sequences: # item = (idx, length); weight_pos=1 -> length at item[1] + seq_len = item[1] + if cur and cur_len + seq_len > packing_length: + packs.append(cur) + cur, cur_len = [], 0 + cur.append(item) + cur_len += seq_len + if cur_len >= packing_length: + packs.append(cur) + cur, cur_len = [], 0 + if is_finished: + if cur: + packs.append(cur) + return packs, [] + return packs, cur + # default: best-fit-decreasing bin packing (https://arxiv.org/pdf/2404.10830) + import binpacking + sequences = binpacking.to_constant_volume(sequences, packing_length, weight_pos=1) + if sequences and not is_finished: + sequences, ret_sequences = sequences[:-1], sequences[-1] + else: + ret_sequences = [] + return sequences, ret_sequences + + +class PackingDataset(Dataset): + PACKING_BATCH_SIZE = 1000 + + def __init__( + self, + template, + dataset, + num_proc: int = 1, + *, + strict: bool = False, + load_from_cache_file: bool = True, + packing_length: Optional[int] = None, + packing_num_proc: int = 1, + packing_strategy: str = 'binpack', + **kwargs, + ): + template.packing = True + template.padding_free = True # TODO: remove + self.template = template + self.dataset = dataset + self.num_proc = num_proc + self.strict = strict + self.load_from_cache_file = load_from_cache_file + self.packing_strategy = packing_strategy + self.packing_length = packing_length or self.template.max_length + self.packing_num_proc = min(packing_num_proc, math.ceil(len(dataset) / self.PACKING_BATCH_SIZE)) + self._out_queue = mp.Queue() + if is_master(): + lengths = self.dataset['lengths'] + offset = 0 + chunked_lengths = split_list(lengths, self.packing_num_proc) + for i in range(self.packing_num_proc): + worker = mp.Process( + target=self.create_packed_idx, args=( + i, + offset, + chunked_lengths[i], + ), daemon=True) + worker.start() + offset += len(chunked_lengths[i]) + self.packed_idx = [[] for _ in range(self.packing_num_proc)] + self.packed_length = [[] for _ in range(self.packing_num_proc)] + desc = 'Packing: ' if self.packing_num_proc == 1 else f'Packing (num_proc={self.packing_num_proc}): ' + with tqdm(total=len(lengths), dynamic_ncols=True, desc=desc) as prog_bar: + finished_workers = 0 + while finished_workers < self.packing_num_proc: + rank, sequences, data_len = self._out_queue.get() + if data_len == -1: + finished_workers += 1 + continue + prog_bar.update(data_len) + self.packed_idx[rank] += [[x[0] for x in seq] for seq in sequences] + self.packed_length[rank] += [sum(x[1] for x in seq) for seq in sequences] + self.packed_idx = list(chain.from_iterable(self.packed_idx)) + self.packed_length = list(chain.from_iterable(self.packed_length)) + else: + self.packed_idx, self.packed_length = None, None + if dist.is_initialized() and is_dist(): + obj_list = [(self.packed_idx, self.packed_length)] + dist.broadcast_object_list(obj_list) + self.packed_idx, self.packed_length = obj_list[0] + + def create_packed_idx(self, rank, offset, lengths): + data = [(i + offset, sum(length) if isinstance(length, list) else length) for i, length in enumerate(lengths)] + i = 0 + input_data = [] + while True: + new_data = data[i:i + self.PACKING_BATCH_SIZE] + input_data += new_data + if not input_data: + break + i += self.PACKING_BATCH_SIZE + is_finished = i >= len(data) + sequences, input_data = calculate_matched_group( + input_data, self.packing_length, is_finished=is_finished, strategy=self.packing_strategy) + self._out_queue.put((rank, sequences, len(new_data))) + self._out_queue.put((rank, [], -1)) + + def __getitem__(self, index): + sequence = self.packed_idx[index] + row = [self.dataset[i] for i in sequence] + return row + + def __len__(self): + return len(self.packed_idx) + + +class IterablePackingDataset(IterableDataset): + + def __init__( + self, + template, + dataset, + num_proc: int = 1, + *, + packing_interval: int = 128, + packing_length: Optional[int] = None, + strict: bool = False, + cyclic: bool = False, + packing_strategy: str = 'binpack', + **kwargs, + ): + template.packing = True + template.padding_free = True # TODO: remove + self.template = template + self.dataset = dataset + self.num_proc = num_proc + self.strict = strict + self.packing_length = packing_length or self.template.max_length + + self.packing_interval = packing_interval + self._in_queue = mp.Queue() + self._out_queue = mp.Queue() + self.workers = [] + self.cyclic = cyclic + self.packing_strategy = packing_strategy + for _ in range(self.num_proc): + worker = mp.Process(target=self._processor, daemon=True) + worker.start() + self.workers.append(worker) + + def _processor(self): + while True: + i, data = self._in_queue.get() + encoded_data = {} + try: + encoded_data = self.template.encode(data, return_length=True) + except Exception as e: + if self.strict and not isinstance(e, MaxLengthError): + raise + self._out_queue.put((i, encoded_data)) + + def _put_data_in_queue(self, iterator) -> int: + for i in range(self.packing_interval): + try: + data = next(iterator) + except StopIteration: + return i + self._in_queue.put((i, data)) + return i + 1 + + def _fetch_data_out_queue(self, last_res, num_samples): + res = [None] * num_samples + for _ in range(num_samples): + i, data = self._out_queue.get() + if not data: + continue + res[i] = data if isinstance(data, list) else [data] + res = [(item, len(item['input_ids'])) for group in res if group for item in group] + last_res += res + return last_res + + @staticmethod + def cyclic_iter(iterable): + while True: + for x in iterable: + yield x + + def set_epoch(self, epoch: int): + if hasattr(self.dataset, 'set_epoch'): + self.dataset.set_epoch(epoch) + + def __iter__(self): + try: + next(iter(self.dataset)) + except StopIteration: + return + + if self.cyclic: + iterator = self.cyclic_iter(self.dataset) + else: + iterator = iter(self.dataset) + data = [] + while True: + num_samples = self._put_data_in_queue(iterator) + finished = num_samples != self.packing_interval + data = self._fetch_data_out_queue(data, num_samples) + sequences, data = calculate_matched_group( + data, self.packing_length, is_finished=finished, strategy=self.packing_strategy) + res = [] + for row in sequences: + res.append([r[0] for r in row]) + yield from res + if finished: + break diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..643e58abdb46e1b4ecd6b4df0c29753c1f868ac1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .core import (DATASET_TYPE, AlpacaPreprocessor, AnthropicMessagesPreprocessor, AutoPreprocessor, ClsPreprocessor, + MessagesPreprocessor, OpenAIMessagesPreprocessor, ResponsePreprocessor, RowPreprocessor) +from .extra import ClsGenerationPreprocessor, GroundingMixin, TextGenerationPreprocessor diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/core.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/core.py new file mode 100644 index 0000000000000000000000000000000000000000..3333b2bda55163331e500c84a44241452dc90884 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/core.py @@ -0,0 +1,705 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import ast +import datasets +import json +import numpy as np +import os +from collections import Counter +from contextlib import contextmanager +from datasets import Dataset as HfDataset +from datasets import Image +from datasets import IterableDataset as HfIterableDataset +from datasets import Sequence, Value +from itertools import chain +from modelscope.hub.utils.utils import get_cache_dir +from packaging import version +from typing import Any, Callable, Dict, List, Literal, Optional, Union + +from swift.template import history_to_messages +from swift.template.template_inputs import normalize_openai_tool_calls +from swift.utils import get_logger, is_dist, is_master, safe_ddp_context + +DATASET_TYPE = Union[HfDataset, HfIterableDataset] + +logger = get_logger() + +_pair_keys = ['messages', 'images', 'videos', 'audios', 'tools', 'objects'] + + +class RowPreprocessor: + standard_keys = _pair_keys + list( + chain.from_iterable([f'{prefix}_{k}' for k in _pair_keys] + for prefix in ['rejected', 'positive', 'negative'])) + [ + 'rejected_response', + 'label', + 'channel', + 'margin', + 'teacher_prompt', + 'chat_template_kwargs', + # Qwen3-TTS + 'ref_audios', + 'audio_codes', + ] + + def __init__(self, + *, + columns: Optional[Dict[str, str]] = None, + dataset_sample: Optional[int] = None, + random_state: Optional[Union[np.random.RandomState, int]] = 42, + traceback_limit: int = 10) -> None: + self.columns = columns or {} + self.origin_columns = self.columns.copy() # Higher priority and raise Error + images_keys = ['images', 'image'] + audios_keys = ['audios', 'audio'] + videos_keys = ['videos', 'video'] + for mm_type in ['images', 'audios', 'videos']: + keys = locals()[f'{mm_type}_keys'] + for key in keys: + self.columns[key] = mm_type + + self.traceback_limit = traceback_limit + self._traceback_counter = 0 + self.dataset_sample = dataset_sample + self.datasets_4 = version.parse(datasets.__version__) >= version.parse('4.0') + if not isinstance(random_state, np.random.RandomState): + random_state = np.random.RandomState(random_state) + self.random_state = random_state + + @staticmethod + def _check_messages(row: Dict[str, Any]) -> None: + if 'messages' not in row: + return + messages = row['messages'] + assert len(messages) > 0, f'messages: {messages}' + # fix swift/SlimOrca (concat) + for message in messages: + keys = set(message.keys()) - {'role', 'content', 'loss', 'loss_scale'} + for key in keys: + message.pop(key) + + for message in messages: + role, content = message['role'], message['content'] + # The terms 'tool' and 'tool_response' have the same meaning, ensuring compatibility. + assert role in {'system', 'user', 'tool_call', 'tool_response', 'tool', 'assistant'}, f'message: {message}' + assert content is not None, f'message: {message}' + + @staticmethod + def _cast_mm_data(row: Dict[str, Any]) -> None: + for key in ['images', 'rejected_images']: + images = row.get(key, None) + if images is None: + continue + + if isinstance(images, str) or (isinstance(images, list) and images and isinstance(images[0], str)): + if isinstance(images, str): + images = [images] + for i, image in enumerate(images): + images[i] = {'bytes': None, 'path': image} + row[key] = images + elif isinstance(images, dict): + row[key] = [images] + + for key in ['videos', 'audios']: + mm_data = row.get(key) + if mm_data is None: + continue + elif isinstance(mm_data, str): + row[key] = [mm_data] + + @staticmethod + def _check_rejected_response(row: Dict[str, Any]) -> None: + if 'rejected_response' in row: + messages = row['messages'] + rejected_response = row['rejected_response'] + if (rejected_response is None + or isinstance(rejected_response, str) and rejected_response == messages[-1]['content']): + raise ValueError(f'rejected_response: {rejected_response}') + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + return row + + def prepare_dataset(self, dataset: DATASET_TYPE) -> DATASET_TYPE: + return dataset + + @staticmethod + def batched_to_rows(batched_row: Dict[str, Any]): + keys = list(batched_row.keys()) + batch_size = len(batched_row[keys[0]]) + return [{key: batched_row[key][i] for key in keys} for i in range(batch_size)] + + @staticmethod + def rows_to_batched(rows: List[Dict[str, Any]]): + batched = {} + for i, row in enumerate(rows): + for k, v in row.items(): + if k not in batched: + batched[k] = [None] * i + batched[k].append(v) + # Make all the lengths of v the same. + for k in set(batched.keys()) - set(row.keys()): + batched[k].append(None) + return batched + + @staticmethod + def _remove_prefix_keys(row, prefix: str): + for k in list(row.keys()): + if k.startswith(prefix): + new_k = k[len(prefix):] + new_v = row.pop(k) + if new_k not in row: + row[new_k] = new_v + + @staticmethod + def _check_objects(row): + objects = row.get('objects') + if objects is None: + return + new_objects = {} + # Ensure the order + for k in ['ref', 'bbox', 'bbox_type', 'image_id']: + if k in objects.keys(): + new_objects[k] = objects[k] + row['objects'] = new_objects + bbox = new_objects['bbox'] + + # check bbox + for box in bbox: + assert len(box) in {2, 4}, f'len(box): {len(box)}' + if len(box) == 2: + continue + if box[0] > box[2]: + box[0], box[2] = box[2], box[0] + if box[1] > box[3]: + box[1], box[3] = box[3], box[1] + + def batched_preprocess(self, batched_row: Dict[str, Any], *, strict: bool, + ignore_max_length_error: bool) -> Dict[str, Any]: + from swift.template import MaxLengthError + batched_row = dict(batched_row) + assert len(batched_row) > 0 + self._remove_prefix_keys(batched_row, '__@') # compat streaming + rows = self.batched_to_rows(batched_row) + + new_rows = [] + for row in rows: + try: + row = self.preprocess(row) + # support [row1, row2, ...] + if row is None: + row = [] + if isinstance(row, dict): + row = [row] + for r in row: + self._check_objects(r) + self._check_rejected_response(r) + self._check_messages(r) + self._cast_mm_data(r) + except Exception as e: + if strict: + logger.warning('To avoid errors, you can pass `strict=False`.') + raise + if isinstance(e, MaxLengthError) and ignore_max_length_error: + pass + elif self.traceback_limit is not None and self._traceback_counter < self.traceback_limit: + import traceback + logger.info(traceback.format_exc()) + logger.warning('👆👆👆There are errors in the dataset, the data will be deleted') + self._traceback_counter += 1 + row = [] + new_rows += row + res = self.rows_to_batched(new_rows) + self._remove_prefix_keys(res, '__#') # compat GRPO + if len(res) == 0: + res['messages'] = [] + + return res + + @staticmethod + def get_features_dataset(dataset: DATASET_TYPE) -> DATASET_TYPE: + if dataset.features is None: + assert isinstance(dataset, HfIterableDataset) + dataset = dataset._resolve_features() + return dataset + + @staticmethod + def safe_rename_columns(dataset, columns): + dataset = RowPreprocessor.get_features_dataset(dataset) + columns_keys = {k.lower(): k for k in dataset.features.keys()} # lower -> lower/upper + safe_columns = {columns_keys[k.lower()]: v for k, v in columns.items() if k.lower() in columns_keys} + + counter = Counter(safe_columns.values()) + for k, new_k in list(safe_columns.items()): + if counter[new_k] > 1: + # For example, if "response" and "answer" match, then no processing is done. + safe_columns.pop(k) + continue + + # e.g. Keep {'query': 'query'} to ensure that the query has the highest priority. + safe_columns = {k: v for k, v in safe_columns.items() if k != v} + if safe_columns: + dataset = dataset.rename_columns(safe_columns) + + return dataset + + @staticmethod + def remove_useless_columns(dataset: DATASET_TYPE) -> DATASET_TYPE: + dataset = RowPreprocessor.get_features_dataset(dataset) + features = dataset.features + k_list = [k for k in RowPreprocessor.standard_keys if k in features] + if len(k_list) != len(features): + dataset = dataset.select_columns(k_list) + return dataset + + @contextmanager + def _patch_arrow_writer(self): + # fix AI-ModelScope/ms_agent_for_agentfabric:all + from datasets.arrow_writer import ArrowWriter + + def _new_init(_self, schema=None, features=None, *args, **kwargs): + + if features is not None: + + if self.datasets_4: + from datasets.features import Json, List + messages_feature = List(Json()) + for key in ['messages', 'rejected_messages', 'positive_messages', 'negative_messages']: + features[key] = messages_feature + features['images'] = List({'bytes': Value(dtype='binary'), 'path': Value(dtype='string')}) + features['objects'] = Json() + features['chat_template_kwargs'] = Json() + else: + messages_feature = [{ + 'role': Value(dtype='string'), + 'content': Value(dtype='string'), + }] + messages_feature_with_loss = [{ + 'role': Value(dtype='string'), + 'content': Value(dtype='string'), + 'loss': Value(dtype='bool'), + 'loss_scale': Value(dtype='float64'), + }] + features['messages'] = messages_feature_with_loss + features['rejected_messages'] = messages_feature_with_loss + features['positive_messages'] = messages_feature + features['negative_messages'] = messages_feature + features['images'] = [{'bytes': Value(dtype='binary'), 'path': Value(dtype='string')}] + features['objects'] = { + 'ref': Sequence(feature=Value(dtype='string'), length=-1), + 'bbox': Sequence(feature=Sequence(feature=Value(dtype='float64'), length=-1), length=-1), + 'bbox_type': Value(dtype='string'), + 'image_id': Sequence(feature=Value(dtype='int64'), length=-1), + } + ArrowWriter.__origin_init__(_self, schema, features, *args, **kwargs) + + ArrowWriter.__origin_init__ = ArrowWriter.__init__ + ArrowWriter.__init__ = _new_init + try: + yield + finally: + ArrowWriter.__init__ = ArrowWriter.__origin_init__ + del ArrowWriter.__origin_init__ + + def _cast_pil_image(self, dataset): + features = dataset.features + for col in ['images', 'rejected_images']: + if (col in features and isinstance(features[col], Image) and getattr(features[col], 'decode', False)): + dataset = dataset.cast_column(col, Image(decode=False)) + return dataset + + def __call__( + self, + dataset: DATASET_TYPE, + *, + num_proc: int = 1, + load_from_cache_file: bool = True, + strict: bool = False, + batch_size: Optional[int] = None, + enable_auto_mapping: bool = False, + ) -> DATASET_TYPE: + from ..utils import sample_dataset + if batch_size is None: + batch_size = 1000 if isinstance(dataset, HfDataset) else 16 + if self.dataset_sample is not None: + dataset = sample_dataset(dataset, self.dataset_sample, True, self.random_state) + + map_kwargs = {'batched': True, 'batch_size': batch_size} + if isinstance(dataset, HfDataset): + if not load_from_cache_file and is_dist() and not is_master(): + load_from_cache_file = True + map_kwargs.update({ + 'num_proc': num_proc, + 'load_from_cache_file': load_from_cache_file, + }) + # compat GRPO: The solution field will be retained. + dataset = RowPreprocessor.get_features_dataset(dataset) + if 'solution' in dataset.features: + with safe_ddp_context(None, True): + if isinstance(dataset, HfDataset) and not dataset.cache_files: + map_kwargs['cache_file_name'] = os.path.join(get_cache_dir(), 'datasets', 'map_cache', + f'{dataset._fingerprint}.arrow') + dataset = dataset.map(lambda x: {'__#solution': x['solution']}, **map_kwargs) + map_kwargs.pop('cache_file_name', None) + dataset = self.safe_rename_columns(dataset, self.origin_columns) + if enable_auto_mapping: + dataset = self.safe_rename_columns(dataset, self.columns) + dataset = self.prepare_dataset(dataset) + dataset = self._cast_pil_image(dataset) + if isinstance(dataset, HfIterableDataset): + # fix: https://github.com/huggingface/datasets/issues/6408 + columns = {k: f'__@{k}' for k in RowPreprocessor.standard_keys if k in dataset.features} + if columns: + dataset = dataset.rename_columns(columns) + + ignore_max_length_error = True + with self._patch_arrow_writer(), safe_ddp_context(None, True): + if isinstance(dataset, HfDataset) and not dataset.cache_files: + map_kwargs['cache_file_name'] = os.path.join(get_cache_dir(), 'datasets', 'map_cache', + f'{dataset._fingerprint}.arrow') + dataset_mapped = dataset.map( + self.batched_preprocess, + fn_kwargs={ + 'strict': strict, + 'ignore_max_length_error': ignore_max_length_error, + }, + remove_columns=list(dataset.features.keys()), + **map_kwargs) + if isinstance(dataset_mapped, HfDataset) and len(dataset) != len(dataset_mapped): + logger.info( + f'Dataset filtered, origin length: {len(dataset)}, filtered dataset length: {len(dataset_mapped)}') + + return dataset_mapped + + +class ResponsePreprocessor(RowPreprocessor): + """Dataset compatible with older versions of ms-swift""" + + def __init__(self, *, columns: Optional[Dict[str, str]] = None, **kwargs) -> None: + super().__init__(columns=columns, **kwargs) + system_keys = ['system', 'system_prompt'] + query_keys = ['query', 'prompt', 'input', 'instruction', 'question', 'problem'] + response_keys = ['response', 'answer', 'output', 'targets', 'target', 'answer_key', 'answers', 'solution' + ] + ['text', 'completion', 'content'] + for key in system_keys: + self.columns[key] = 'system' + for key in query_keys: + self.columns[key] = 'query' + for key in response_keys: + self.columns[key] = 'response' + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + response = row.pop('response', None) + if response is not None: + if isinstance(response, (list, tuple)): + from transformers.utils import strtobool + + # sometimes response is a list, pick one randomly + if strtobool(os.environ.get('RANDOM_DATASET_RESPONSE', 'False')): + response = self.random_state.choice(response) + else: + response = response[0] + history = row.pop('history', None) or [] + query = row.pop('query', None) + system = row.pop('system', None) + if isinstance(history, str): # e.g. "[['query1', 'response1']]" + history = ast.literal_eval(history) + history.append([query, response]) + + row.update({'messages': history_to_messages(history, system)}) + return row + + +class AlpacaPreprocessor(ResponsePreprocessor): + + @classmethod + def concat_inst_input(cls, instruction, input_): + if instruction and input_: + query = f'{instruction}\n{input_}' + else: + query = instruction or input_ + assert isinstance(query, str), f'query: {query}' + return query + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + instruction = row.pop('instruction', None) + input_ = row.pop('input', None) + output = row.pop('output', None) + if output is not None: + row['response'] = output + row['query'] = self.concat_inst_input(instruction, input_) + return super().preprocess(row) + + +def default_repair_messages(s: Union[str, Any]) -> Any: + if isinstance(s, str): + return ast.literal_eval(s) + return s + + +class MessagesPreprocessor(RowPreprocessor): + + def __init__( + self, + *, + # If set to None, automatic matching will be performed. + role_key: Optional[str] = None, # 'role', 'from' + content_key: Optional[str] = None, # 'content', 'value' + user_role: Optional[str] = None, # 'user', 'human' + assistant_role: Optional[str] = None, # 'assistant', 'gpt', 'bot' + system_role: str = 'system', + # 'conversation', 'conversations' -> 'messages' + columns: Optional[Dict[str, str]] = None, + repair_messages: Callable[[Union[str, List[Dict[str, str]]]], + Optional[List[Dict[str, str]]]] = default_repair_messages, + inner_key: Optional[str] = None, + message_format: Literal['auto', 'swift', 'openai', 'anthropic'] = 'auto', + **kwargs): + super().__init__(columns=columns, **kwargs) + self.role_keys = ['role', 'from'] if role_key is None else [role_key] + self.content_keys = ['content', 'value'] if content_key is None else [content_key] + self.user_roles = ['user', 'human'] if user_role is None else [user_role] + self.assistant_roles = ['assistant', 'gpt', 'bot'] if assistant_role is None else [assistant_role] + self.tool_call_roles = ['function_call'] + self.tool_response_roles = ['function_response', 'observation', 'observations'] + + self.system_role = system_role + self.repair_messages = repair_messages + self.inner_key = inner_key + self.message_format = message_format + + message_keys = ['messages', 'conversation', 'conversations'] + for key in message_keys: + self.columns[key] = 'messages' + # sharegptq + system_keys = ['system', 'system_prompt'] + if system_role not in system_keys: + system_keys.append(system_role) + for key in system_keys: + self.columns[key] = 'system' + + @staticmethod + def _is_sharegpt_format(message: Dict[str, str]) -> bool: + if 'role' in message or 'content' in message: + return False + return True + + def sharegpt_to_messages(self, messages: List[Dict[str, str]], system: Optional[str]) -> List[Dict[str, str]]: + self._to_std_key(messages, 'user', self.user_roles) + self._to_std_key(messages, 'assistant', self.assistant_roles) + new_messages = [] + if system is not None: + new_messages.append({'role': 'system', 'content': system}) + for message in messages: + user_message = {'role': 'user', 'content': message['user']} + assistant_message = {'role': 'assistant', 'content': message['assistant']} + new_messages.append(user_message) + new_messages.append(assistant_message) + return new_messages + + def to_std_messages(self, messages: List[Dict[str, str]], system: Optional[str]) -> None: + if messages[0]['role'] == self.system_role: + messages[0]['role'] = 'system' + elif system is not None: + messages.insert(0, {'role': 'system', 'content': system}) + for message in messages: + role = message['role'] + if role in self.user_roles: + message['role'] = 'user' + elif role in self.assistant_roles: + message['role'] = 'assistant' + elif role.replace('-', '_') in self.tool_call_roles: + message['role'] = 'tool_call' + elif role.replace('-', '_') in self.tool_response_roles: + message['role'] = 'tool_response' + + @staticmethod + def openai_to_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Convert OpenAI tool-call messages to the SWIFT canonical roles.""" + return normalize_openai_tool_calls(messages) + + @staticmethod + def _anthropic_image_source(block: Dict[str, Any]) -> str: + source = block.get('source') or {} + source_type = source.get('type') + if source_type == 'base64': + media_type = source.get('media_type') + data = source.get('data') + if not media_type or not data: + raise ValueError(f'Invalid Anthropic base64 image block: {block}') + return f'data:{media_type};base64,{data}' + if source_type == 'url': + url = source.get('url') + if not url: + raise ValueError(f'Invalid Anthropic URL image block: {block}') + return url + raise ValueError(f'Unsupported Anthropic image source type: {source_type}') + + @classmethod + def _anthropic_block_content(cls, content: Any, media: Dict[str, List[Any]]) -> Any: + if not isinstance(content, list): + return content + parts = [] + for block in content: + block_type = block.get('type') + if block_type == 'text': + parts.append(block.get('text', '')) + elif block_type == 'image': + parts.append('') + media['images'].append(cls._anthropic_image_source(block)) + else: + raise ValueError(f'Unsupported Anthropic content block type: {block_type}') + return ''.join(parts) + + @classmethod + def anthropic_to_messages(cls, + messages: List[Dict[str, Any]], + media: Optional[Dict[str, List[Any]]] = None) -> List[Dict[str, Any]]: + """Convert Anthropic content blocks to the SWIFT canonical roles.""" + media = media if media is not None else {'images': []} + new_messages = [] + for message in messages: + content = message.get('content') + if not isinstance(content, list): + new_messages.append(message) + continue + + pending_content = [] + message_metadata = {key: message[key] for key in ['loss', 'loss_scale'] if key in message} + + def flush_content(): + if pending_content: + new_messages.append({ + 'role': message['role'], + 'content': ''.join(pending_content), + **message_metadata + }) + pending_content.clear() + + for block in content: + block_type = block.get('type') + if block_type == 'text': + pending_content.append(block.get('text', '')) + elif block_type == 'image': + pending_content.append('') + media['images'].append(cls._anthropic_image_source(block)) + elif block_type == 'tool_use': + flush_content() + new_messages.append({ + 'role': 'tool_call', + 'content': { + 'name': block['name'], + 'arguments': block.get('input', {}), + }, + **message_metadata, + }) + elif block_type == 'tool_result': + flush_content() + new_messages.append({ + 'role': 'tool_response', + 'content': cls._anthropic_block_content(block.get('content', ''), media), + **message_metadata, + }) + else: + raise ValueError(f'Unsupported Anthropic content block type: {block_type}') + flush_content() + return new_messages + + def normalize_provider_messages(self, messages: List[Dict[str, Any]], + media: Dict[str, List[Any]]) -> List[Dict[str, Any]]: + message_format = self.message_format + if message_format == 'auto': + if any(message.get('tool_calls') for message in messages): + message_format = 'openai' + elif any( + isinstance(message.get('content'), list) and any( + block.get('type') in {'tool_use', 'tool_result'} + or block.get('type') == 'image' and 'source' in block for block in message['content']) + for message in messages): + message_format = 'anthropic' + else: + message_format = 'swift' + if message_format == 'openai': + return self.openai_to_messages(messages) + if message_format == 'anthropic': + return self.anthropic_to_messages(messages, media) + return messages + + @staticmethod + def _to_std_key(messages: List[Dict[str, str]], std_key: str, optional_keys: List[str]) -> None: + for message in messages: + for key in optional_keys: + if key in message: + message[std_key] = message.pop(key) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + if 'rejected_messages' in row: + rejected = MessagesPreprocessor.preprocess(self, {'messages': row['rejected_messages']}) + row['rejected_messages'] = rejected['messages'] if rejected else None + messages = row['messages'] + if self.inner_key is not None: + messages = messages[self.inner_key] + messages: Optional[List[Dict[str, str]]] = self.repair_messages(messages) + if not messages or isinstance(messages, str): + return + media = {'images': []} + messages = self.normalize_provider_messages(messages, media) + for key, value in media.items(): + if value: + assert not row.get(key), f'Cannot mix Anthropic content blocks with the top-level `{key}` field.' + row[key] = value + self._to_std_key(messages, 'role', self.role_keys) + self._to_std_key(messages, 'content', self.content_keys) + system = row.pop('system', None) + if self._is_sharegpt_format(messages[0]): + messages = self.sharegpt_to_messages(messages, system) + else: + self.to_std_messages(messages, system) # inplace + row['messages'] = messages + return row + + +class OpenAIMessagesPreprocessor(MessagesPreprocessor): + + def __init__(self, **kwargs): + super().__init__(message_format='openai', **kwargs) + + +class AnthropicMessagesPreprocessor(MessagesPreprocessor): + + def __init__(self, **kwargs): + super().__init__(message_format='anthropic', **kwargs) + + +class ClsPreprocessor(ResponsePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + res = super().preprocess(row) + res['label'] = int(res['label']) + return res + + +class AutoPreprocessor: + + def __init__(self, *, columns: Optional[Dict[str, str]] = None, **kwargs) -> None: + self.columns = columns or {} + self.kwargs = kwargs + + def _get_preprocessor(self, dataset: DATASET_TYPE) -> RowPreprocessor: + features = dataset.features + for key in ['conversation', 'conversations', 'messages']: + if key in features: + return MessagesPreprocessor(**self.kwargs) + if 'instruction' in features and 'input' in features: + return AlpacaPreprocessor(**self.kwargs) + return ResponsePreprocessor(**self.kwargs) + + def __call__( + self, + dataset: DATASET_TYPE, + *, + num_proc: int = 1, + load_from_cache_file: bool = True, + **kwargs, + ) -> DATASET_TYPE: + dataset = RowPreprocessor.safe_rename_columns(dataset, self.columns) + preprocessor = self._get_preprocessor(dataset) + return preprocessor(dataset, num_proc=num_proc, load_from_cache_file=load_from_cache_file, **kwargs) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/extra.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/extra.py new file mode 100644 index 0000000000000000000000000000000000000000..d53417d143b909caf764d0ed2091f985676279e9 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/preprocessor/extra.py @@ -0,0 +1,111 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import numpy as np +from typing import Any, Dict, List, Optional + +from .core import ResponsePreprocessor + + +class GroundingMixin: + """This class offers prompts to the grounding task""" + task_type: Optional[str] = None + + _grounding_language_mixin = [0.8, 0.2] + _grounding_prompts = { + 'grounding': { + 'en': [('', ''), ('The positions of is', ''), + ('Find the positions of ', ''), ('Where is ', ''), + ('Find ', ''), ('Show me ', ''), + ('Detect ', ''), ('Locate ', ''), + ('Tell me the location of ', ''), ('Give the location of ', ''), + ('Provide the bounding box coordinate of ', '')], + 'zh': [('', ''), ('的位置在图片中', ''), ('在图片中', ''), + ('在', ''), ('找到的位置', ''), ('在哪里', ''), + ('提供的坐标位置', '')] + }, + 'caption': { + 'en': [ + ('', ''), + ('The object at position ', ''), + ('This is', ''), + ('What is the object at ', ''), + ('Describe ', ''), + (' is', ''), + ('The bounding box coordinate contains', ''), + ], + 'zh': [ + ('', ''), + ('是什么', ''), + ('的位置包含', ''), + ('描述', ''), + ('中是', ''), + ('坐标描述了什么', ''), + ('描述中的事物', ''), + ] + }, + } + + def construct_grounding_prompt(self): + # TODO Only support one bbox to one object + lang = np.random.choice(['en', 'zh'], p=[0.8, 0.2]) + prompts = GroundingMixin._grounding_prompts[self.task_type][lang] + query, response = prompts[np.random.choice(range(len(prompts)))] + return query, response + + +class TextGenerationPreprocessor(ResponsePreprocessor): + + def __init__(self, + *, + prompt: str, + query_tag: str = '{{QUERY}}', + columns: Optional[Dict[str, str]] = None, + **kwargs) -> None: + self.query_tag = query_tag + self.prompt = prompt + super().__init__(columns=columns, **kwargs) + + def preprocess(self, row: Dict[str, Any]) -> Dict[str, Any]: + row['query'] = self.prompt.replace(self.query_tag, row['query']) + return super().preprocess(row) + + +class ClsGenerationPreprocessor(ResponsePreprocessor): + + def __init__(self, + labels: List[str], + *, + task: str, + is_pair_seq: bool = False, + columns: Optional[Dict[str, str]] = None, + **kwargs) -> None: + self.labels = labels + self.task = task + self.is_pair_seq = is_pair_seq + + category = ', '.join(labels) + self.sentence2_key = 'sentence2' + self.label_key = 'label' + if is_pair_seq: + self.sentence_key = 'sentence1' + inputs = 'Sentence1: {sentence1}\nSentence2: {sentence2}' + else: + self.sentence_key = 'sentence' + inputs = 'Sentence: {sentence}' + self.prompt = f"""Task: {task} +{inputs} +Category: {category} +Output:""" + super().__init__(columns=columns, **kwargs) + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + label = row.pop(self.label_key, None) + if label is None: + return + + if self.is_pair_seq: + query = self.prompt.format(sentence1=row.pop(self.sentence_key), sentence2=row.pop(self.sentence2_key)) + else: + query = self.prompt.format(sentence=row.pop(self.sentence_key)) + row['query'] = query + row['response'] = self.labels[int(label)] + return super().preprocess(row) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/register.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/register.py new file mode 100644 index 0000000000000000000000000000000000000000..c856c807742fc05bdfc578ec23e849e51c22a363 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/register.py @@ -0,0 +1,115 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import os +from copy import deepcopy +from typing import Any, Dict, List, Optional, Union + +from swift.utils import get_logger, use_hf_hub +from .dataset_meta import DATASET_MAPPING, DatasetMeta, SubsetDataset +from .preprocessor import AutoPreprocessor, MessagesPreprocessor + +logger = get_logger() + + +def get_dataset_list(): + datasets = [] + for key in DATASET_MAPPING: + if use_hf_hub(): + if key[1]: + datasets.append(key[1]) + else: + if key[0]: + datasets.append(key[0]) + return datasets + + +def register_dataset(dataset_meta: DatasetMeta, *, exist_ok: bool = False) -> None: + """Register dataset + + Args: + dataset_meta: The `DatasetMeta` info of the dataset. + exist_ok: If the dataset id exists, raise error or update it. + """ + if dataset_meta.dataset_name: + dataset_name = dataset_meta.dataset_name + else: + dataset_name = dataset_meta.ms_dataset_id, dataset_meta.hf_dataset_id, dataset_meta.dataset_path + if not exist_ok and dataset_name in DATASET_MAPPING: + raise ValueError(f'The `{dataset_name}` has already been registered in the DATASET_MAPPING.') + + DATASET_MAPPING[dataset_name] = dataset_meta + + +def _preprocess_d_info(d_info: Dict[str, Any], *, base_dir: Optional[str] = None) -> Dict[str, Any]: + d_info = deepcopy(d_info) + + columns = None + if 'columns' in d_info: + columns = d_info.pop('columns') + + if 'messages' in d_info: + d_info['preprocess_func'] = MessagesPreprocessor(**d_info.pop('messages'), columns=columns) + else: + d_info['preprocess_func'] = AutoPreprocessor(columns=columns) + + if 'dataset_path' in d_info: + dataset_path = d_info.pop('dataset_path') + if base_dir is not None and not os.path.isabs(dataset_path): + dataset_path = os.path.join(base_dir, dataset_path) + dataset_path = os.path.abspath(os.path.expanduser(dataset_path)) + + d_info['dataset_path'] = dataset_path + + if 'subsets' in d_info: + subsets = d_info.pop('subsets') + for i, subset in enumerate(subsets): + if isinstance(subset, dict): + subsets[i] = SubsetDataset(**_preprocess_d_info(subset)) + d_info['subsets'] = subsets + return d_info + + +def _register_d_info(d_info: Dict[str, Any], *, base_dir: Optional[str] = None) -> DatasetMeta: + """Register a single dataset to dataset mapping + + Args: + d_info: The dataset info + """ + d_info = _preprocess_d_info(d_info, base_dir=base_dir) + dataset_meta = DatasetMeta(**d_info) + register_dataset(dataset_meta) + return dataset_meta + + +def register_dataset_info(dataset_info: Union[str, List[str], None] = None) -> List[DatasetMeta]: + """Register dataset from the `dataset_info.json` or a custom dataset info file + This is used to deal with the datasets defined in the json info file. + + Args: + dataset_info: The dataset info path + """ + # dataset_info_path: path, json or None + if dataset_info is None: + dataset_info = os.path.join(os.path.dirname(__file__), 'data', 'dataset_info.json') + assert isinstance(dataset_info, (str, list)) + base_dir = None + log_msg = None + if isinstance(dataset_info, str): + dataset_path = os.path.abspath(os.path.expanduser(dataset_info)) + if os.path.isfile(dataset_path): + log_msg = dataset_path + base_dir = os.path.dirname(dataset_path) + with open(dataset_path, 'r', encoding='utf-8') as f: + dataset_info = json.load(f) + else: + dataset_info = json.loads(dataset_info) # json + if len(dataset_info) == 0: + return [] + res = [] + for d_info in dataset_info: + res.append(_register_d_info(d_info, base_dir=base_dir)) + + if log_msg is None: + log_msg = dataset_info if len(dataset_info) < 5 else list(dataset_info.keys()) + logger.info(f'Successfully registered `{log_msg}`.') + return res diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/dataset/utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e78abf12f275f1d74d81a617ba984121fe9a2e2a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/dataset/utils.py @@ -0,0 +1,153 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import inspect +import numpy as np +import os +import tempfile +from datasets import Dataset as HfDataset +from modelscope.hub.utils.utils import get_cache_dir +from torch.utils.data import Dataset +from typing import Any, Callable, Dict, Optional, Union + +from swift.template import MaxLengthError, Template +from swift.utils import get_logger +from .preprocessor import RowPreprocessor + +logger = get_logger() + + +def sample_dataset( + dataset: HfDataset, + dataset_sample: Optional[int], + shuffle: bool = True, + random_state: Optional[np.random.RandomState] = None, + shuffle_all: bool = False, # For compatibility, this defaults to False. +) -> HfDataset: + """Sample dataset by a dataset_sample number + Args: + dataset: The dataset instance, iterable dataset is not supported + dataset_sample: The sample number + shuffle: Whether to perform random sampling on non-streaming datasets + random_state: The random state + Returns: + The sampled dataset + """ + if dataset_sample is None: + return dataset + + n_repeat_sample = dataset_sample // len(dataset) + n_remain_sample = dataset_sample % len(dataset) + if n_repeat_sample >= 1 and n_remain_sample >= 1: + logger.warning(f'dataset_sample:{dataset_sample} is greater than len(dataset):{len(dataset)}, ' + 'repeated sampling will be performed.') + idx = np.tile(range(len(dataset)), n_repeat_sample) + if random_state is None: + random_state = np.random.RandomState() + if n_remain_sample >= 1: + if shuffle: + idx_remain = random_state.permutation(len(dataset))[:n_remain_sample] + else: + idx_remain = np.arange(n_remain_sample) + idx = np.concatenate([idx, idx_remain]) + if n_repeat_sample >= 1 and shuffle and shuffle_all: + random_state.shuffle(idx) + dataset = dataset.select(idx) + return dataset + + +class LazyLLMDataset(Dataset): + """This class if used to lazy tokenize the dataset, and skips bad ones when training""" + + def __init__(self, + dataset: HfDataset, + encode_func: Callable[[Dict[str, Any]], Dict[str, Any]], + *, + n_try_fetch: int = 10, + strict: bool = False, + random_state: Optional[Union[np.random.RandomState, int]] = None, + traceback_limit: int = 10) -> None: + self.dataset = dataset + self.encode_func = encode_func + + n_try_fetch = 1 if strict else min(n_try_fetch, len(self.dataset)) + assert n_try_fetch >= 1 + self.strict = strict + self.n_try_fetch = n_try_fetch + + if not isinstance(random_state, np.random.RandomState): + random_state = np.random.RandomState(random_state) + self.random_state = random_state + + self.traceback_limit = traceback_limit + self._traceback_counter = 0 + self._idx = 0 + self._idx_list = self.random_state.permutation(len(self.dataset)).tolist() + + def __getitem__(self, idx: int) -> Dict[str, Any]: + if isinstance(idx, str): + return self.dataset[idx] + for i in range(self.n_try_fetch): + if i > 0: + idx = self._idx_list[self._idx] + self._idx = (self._idx + 1) % len(self.dataset) + data = self.dataset[idx] + try: + return self.encode_func(data, return_length=True) + except Exception as e: + if self.strict: + logger.warning('To avoid errors, you can pass `strict=False`.') + raise + if isinstance(e, MaxLengthError): + continue + if self.traceback_limit is not None and self._traceback_counter < self.traceback_limit: + import traceback + logger.info(traceback.format_exc()) + logger.warning('👆👆👆There are errors in the template.encode, ' + 'and another piece of data will be randomly selected.') + self._traceback_counter += 1 + + raise ValueError('Failed to retrieve the dataset. You can avoid this issue by increasing `max_length` or ' + 'modifying the `truncation_strategy`.') + + def __len__(self) -> int: + return len(self.dataset) + + +class EncodePreprocessor(RowPreprocessor): + + def __init__(self, template: 'Template'): + super().__init__() + self.template = template + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + return self.template.encode(row, return_length=True) + + +class AddLengthPreprocessor(EncodePreprocessor): + + def preprocess(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]: + encoded = super().preprocess(row) + row['lengths'] = encoded['lengths'] + return row + + +TEMP_DIR_POOL = {} + + +def get_temporary_cache_files_directory(prefix=None): + if prefix is None: + import datasets.config + prefix = datasets.config.TEMP_CACHE_DIR_PREFIX + if prefix in TEMP_DIR_POOL: + TEMP_DIR = TEMP_DIR_POOL[prefix] + else: + tmp_dir = os.path.join(get_cache_dir(), 'tmp') + os.makedirs(tmp_dir, exist_ok=True) + kwargs = {} + parameters = inspect.signature(tempfile.TemporaryDirectory.__init__).parameters + if 'ignore_cleanup_errors' in parameters: + kwargs['ignore_cleanup_errors'] = True + TEMP_DIR = tempfile.TemporaryDirectory(prefix=prefix, dir=tmp_dir, **kwargs) + logger.info(f'create tmp_dir: {TEMP_DIR.name}') + TEMP_DIR_POOL[prefix] = TEMP_DIR + + return TEMP_DIR.name diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/hub/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/hub/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..97bd85a3318f25d990a0730d29bb699282a88896 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/hub/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .hub import HFHub, MSHub, get_hub diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/hub/constant.py b/pipelines/Tr006_V2/project/ms-swift/swift/hub/constant.py new file mode 100644 index 0000000000000000000000000000000000000000..53b46b0f18fa7a8cee4f14e43152769b4be01b38 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/hub/constant.py @@ -0,0 +1,6 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from modelscope.hub import constants + +constants.API_HTTP_CLIENT_TIMEOUT = 5 +constants.API_FILE_DOWNLOAD_TIMEOUT = 300 +constants.API_FILE_DOWNLOAD_CHUNK_SIZE = 1024 * 1024 * 16 diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/hub/hub.py b/pipelines/Tr006_V2/project/ms-swift/swift/hub/hub.py new file mode 100644 index 0000000000000000000000000000000000000000..91ca8e32c67fcf59ab8e5c0d226b7717172062cc --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/hub/hub.py @@ -0,0 +1,462 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import huggingface_hub +import logging +import modelscope +import os +import tempfile +from contextlib import contextmanager +from functools import partial +from huggingface_hub import RepoUrl +from huggingface_hub.hf_api import api, future_compatible +from packaging import version +from pathlib import Path +from requests.exceptions import HTTPError +from transformers import trainer +from transformers.utils import strtobool +from typing import List, Literal, Optional, Union + +from swift.utils import get_logger, ms_logger_context, use_hf_hub + +logger = get_logger() + + +class HubOperation: + + @classmethod + @contextmanager + def patch_hub(cls): + yield + + @classmethod + def try_login(cls, token: Optional[str] = None) -> bool: + """Try to login to the hub + + Args: + token: The hub token to use + + Returns: + bool: Whether login is successful + """ + raise NotImplementedError + + @classmethod + def create_model_repo(cls, repo_id: str, token: Optional[str] = None, private: bool = False): + """Create a model repo on the hub + + Args: + repo_id: The model id of the hub + token: The hub token to use + private: If is a private repo + """ + raise NotImplementedError + + @classmethod + def push_to_hub(cls, + repo_id: str, + folder_path: Union[str, Path], + path_in_repo: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + token: Optional[Union[str, bool]] = None, + private: bool = False, + revision: Optional[str] = 'master', + ignore_patterns: Optional[Union[List[str], str]] = None, + **kwargs): + """Push a model-like folder to the hub + + Args: + repo_id: The repo id + folder_path: The local folder path + path_in_repo: Which remote folder to put the local files in + commit_message: The commit message of git + commit_description: The commit description + token: The hub token + private: Private hub or not + revision: The revision to push to + ignore_patterns: The ignore file patterns + """ + raise NotImplementedError + + @classmethod + def load_dataset(cls, + dataset_id: str, + subset_name: str, + split: str, + streaming: bool = False, + revision: Optional[str] = None): + """Load a dataset from the repo + + Args: + dataset_id: The dataset id + subset_name: The subset name of the dataset + split: The split info + streaming: Streaming mode + revision: The revision of the dataset + + Returns: + The Dataset instance + """ + raise NotImplementedError + + @classmethod + def download_model(cls, + model_id_or_path: Optional[str] = None, + revision: Optional[str] = None, + download_model: bool = True, + ignore_patterns: Optional[List[str]] = None, + **kwargs): + """Download model from the hub + + Args: + model_id_or_path: The model id + revision: The model revision + download_model: Whether downloading bin/safetensors files, this is usually useful when only + using tokenizer + ignore_patterns: Custom ignore pattern + **kwargs: + + Returns: + The local dir + """ + raise NotImplementedError + + +class MSHub(HubOperation): + ms_token = None + + @staticmethod + def create_repo(repo_id: str, + *, + token: Optional[Union[str, bool]] = None, + private: bool = False, + **kwargs) -> RepoUrl: + """ + Create a new repository on the hub. + + Args: + repo_id: The ID of the repository to create. + token: The authentication token to use. + private: Whether the repository should be private. + **kwargs: Additional arguments. + + Returns: + RepoUrl: The URL of the created repository. + """ + hub_model_id = MSHub.create_model_repo(repo_id, token, private) + return RepoUrl(url=hub_model_id, ) + + @staticmethod + @future_compatible + def upload_folder( + self, + *, + repo_id: str, + folder_path: Union[str, Path], + path_in_repo: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + token: Optional[Union[str, bool]] = None, + revision: Optional[str] = 'master', + ignore_patterns: Optional[Union[List[str], str]] = None, + **kwargs, + ): + from modelscope.utils.repo_utils import CommitInfo + MSHub.push_to_hub(repo_id, folder_path, path_in_repo, commit_message, commit_description, token, True, revision, + ignore_patterns) + return CommitInfo( + commit_url=f'https://www.modelscope.cn/models/{repo_id}/files', + commit_message=commit_message, + commit_description=commit_description, + oid=None, + ) + + @classmethod + @contextmanager + def patch_hub(cls): + hub_create_repo = huggingface_hub.create_repo + hub_upload_folder = huggingface_hub.upload_folder + trainer_create_repo = getattr(trainer, 'create_repo', None) + trainer_upload_folder = getattr(trainer, 'upload_folder', None) + + huggingface_hub.create_repo = cls.create_repo + huggingface_hub.upload_folder = partial(cls.upload_folder, api) + trainer.create_repo = cls.create_repo + trainer.upload_folder = partial(cls.upload_folder, api) + try: + yield + finally: + huggingface_hub.create_repo = hub_create_repo + huggingface_hub.upload_folder = hub_upload_folder + trainer.create_repo = trainer_create_repo + trainer.upload_folder = trainer_upload_folder + + @classmethod + def try_login(cls, token: Optional[str] = None) -> bool: + from modelscope import HubApi + if token is None: + token = os.environ.get('MODELSCOPE_API_TOKEN') + if token: + api = HubApi() + api.login(token) + return True + return False + + @classmethod + def create_model_repo(cls, repo_id: str, token: Optional[str] = None, private: bool = False) -> str: + from modelscope import HubApi + from modelscope.hub.api import ModelScopeConfig + from modelscope.hub.constants import ModelVisibility + assert repo_id is not None, 'Please enter a valid hub_model_id' + + if not cls.try_login(token): + raise ValueError('Please specify a token by `--hub_token` or `MODELSCOPE_API_TOKEN=xxx`') + cls.ms_token = token + visibility = ModelVisibility.PRIVATE if private else ModelVisibility.PUBLIC + api = HubApi() + if '/' not in repo_id: + user_name = ModelScopeConfig.get_user_info()[0] + assert isinstance(user_name, str) + hub_model_id = f'{user_name}/{repo_id}' + logger.info(f"'/' not in hub_model_id, pushing to personal repo {hub_model_id}") + try: + api.create_model(repo_id, visibility) + except HTTPError: + # The remote repository has been created + pass + + with tempfile.TemporaryDirectory() as temp_cache_dir: + from modelscope.hub.repository import Repository + repo = Repository(temp_cache_dir, repo_id) + cls.add_patterns_to_gitattributes(repo, ['*.safetensors', '*.bin', '*.pt']) + # Add 'runs/' to .gitignore, ignore tensorboard files + cls.add_patterns_to_gitignore(repo, ['runs/', 'images/']) + cls.add_patterns_to_file( + repo, + 'configuration.json', ['{"framework": "pytorch", "task": "text-generation", "allow_remote": true}'], + ignore_push_error=True) + # Add '*.sagemaker' to .gitignore if using SageMaker + if os.environ.get('SM_TRAINING_ENV'): + cls.add_patterns_to_gitignore(repo, ['*.sagemaker-uploading', '*.sagemaker-uploaded'], + 'Add `*.sagemaker` patterns to .gitignore') + return repo_id + + @classmethod + def push_to_hub(cls, + repo_id: str, + folder_path: Union[str, Path], + path_in_repo: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + token: Optional[Union[str, bool]] = None, + private: bool = False, + revision: Optional[str] = 'master', + ignore_patterns: Optional[Union[List[str], str]] = None, + **kwargs): + cls.create_model_repo(repo_id, token, private) + from modelscope import push_to_hub + commit_message = commit_message or 'Upload folder using api' + if commit_description: + commit_message = commit_message + '\n' + commit_description + if not os.path.exists(os.path.join(folder_path, 'configuration.json')): + with open(os.path.join(folder_path, 'configuration.json'), 'w', encoding='utf-8') as f: + f.write('{"framework": "pytorch", "task": "text-generation", "allow_remote": true}') + if ignore_patterns: + ignore_patterns = [p for p in ignore_patterns if p != '_*'] + if path_in_repo: + # We don't support part submit for now + path_in_repo = os.path.basename(folder_path) + folder_path = os.path.dirname(folder_path) + ignore_patterns = [] + if revision is None or revision == 'main': + revision = 'master' + push_to_hub( + repo_id, + folder_path, + token or cls.ms_token, + private, + commit_message=commit_message, + ignore_file_pattern=ignore_patterns, + revision=revision, + tag=path_in_repo) + + @classmethod + def load_dataset(cls, + dataset_id: str, + subset_name: str, + split: str, + streaming: bool = False, + revision: Optional[str] = None, + download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists', + token: Optional[str] = None, + **kwargs): + from modelscope import MsDataset + cls.try_login(token) + if revision is None or revision == 'main': + revision = 'master' + load_kwargs = {} + if version.parse(modelscope.__version__) >= version.parse('1.29.1'): + load_kwargs['trust_remote_code'] = True + with ms_logger_context(logging.ERROR): + return MsDataset.load( + dataset_id, + subset_name=subset_name, + split=split, + version=revision, + download_mode=download_mode, + use_streaming=streaming, + **load_kwargs, + ) + + @classmethod + def download_model(cls, + model_id_or_path: Optional[str] = None, + revision: Optional[str] = None, + ignore_patterns: Optional[List[str]] = None, + token: Optional[str] = None, + **kwargs): + cls.try_login(token) + if revision is None or revision == 'main': + revision = 'master' + logger.info(f'Downloading the model from ModelScope Hub, model_id: {model_id_or_path}') + from modelscope import snapshot_download + return snapshot_download(model_id_or_path, revision, ignore_patterns=ignore_patterns, **kwargs) + + @staticmethod + def add_patterns_to_file(repo, + file_name: str, + patterns: List[str], + commit_message: Optional[str] = None, + ignore_push_error=False) -> None: + if isinstance(patterns, str): + patterns = [patterns] + if commit_message is None: + commit_message = f'Add `{patterns[0]}` patterns to {file_name}' + + # Get current file content + repo_dir = repo.model_dir + file_path = os.path.join(repo_dir, file_name) + if os.path.exists(file_path): + with open(file_path, 'r', encoding='utf-8') as f: + current_content = f.read() + else: + current_content = '' + # Add the patterns to file + content = current_content + for pattern in patterns: + if pattern not in content: + if len(content) > 0 and not content.endswith('\n'): + content += '\n' + content += f'{pattern}\n' + + # Write the file if it has changed + if content != current_content: + with open(file_path, 'w', encoding='utf-8') as f: + logger.debug(f'Writing {file_name} file. Content: {content}') + f.write(content) + try: + repo.push(commit_message) + except Exception as e: + if ignore_push_error: + pass + else: + raise e + + @staticmethod + def add_patterns_to_gitignore(repo, patterns: List[str], commit_message: Optional[str] = None) -> None: + MSHub.add_patterns_to_file(repo, '.gitignore', patterns, commit_message, ignore_push_error=True) + + @staticmethod + def add_patterns_to_gitattributes(repo, patterns: List[str], commit_message: Optional[str] = None) -> None: + new_patterns = [] + suffix = 'filter=lfs diff=lfs merge=lfs -text' + for pattern in patterns: + if suffix not in pattern: + pattern = f'{pattern} {suffix}' + new_patterns.append(pattern) + file_name = '.gitattributes' + if commit_message is None: + commit_message = f'Add `{patterns[0]}` patterns to {file_name}' + MSHub.add_patterns_to_file(repo, file_name, new_patterns, commit_message, ignore_push_error=True) + + +class HFHub(HubOperation): + + @classmethod + def try_login(cls, token: Optional[str] = None) -> bool: + pass + + @classmethod + def create_model_repo(cls, repo_id: str, token: Optional[str] = None, private: bool = False) -> str: + return api.create_repo(repo_id, token=token, private=private) + + @classmethod + def push_to_hub(cls, + repo_id: str, + folder_path: Union[str, Path], + path_in_repo: Optional[str] = None, + commit_message: Optional[str] = None, + commit_description: Optional[str] = None, + token: Optional[Union[str, bool]] = None, + private: bool = False, + revision: Optional[str] = 'master', + ignore_patterns: Optional[Union[List[str], str]] = None, + **kwargs): + cls.create_model_repo(repo_id, token, private) + if revision is None or revision == 'master': + revision = 'main' + return api.upload_folder( + repo_id=repo_id, + folder_path=folder_path, + path_in_repo=path_in_repo, + commit_message=commit_message, + commit_description=commit_description, + token=token, + revision=revision, + ignore_patterns=ignore_patterns, + **kwargs) + + @classmethod + def load_dataset(cls, + dataset_id: str, + subset_name: str, + split: str, + streaming: bool = False, + revision: Optional[str] = None, + download_mode: Literal['force_redownload', 'reuse_dataset_if_exists'] = 'reuse_dataset_if_exists', + num_proc: Optional[int] = None, + **kwargs): + from datasets import load_dataset + if revision is None or revision == 'master': + revision = 'main' + return load_dataset( + dataset_id, + name=subset_name, + split=split, + streaming=streaming, + revision=revision, + download_mode=download_mode, + num_proc=num_proc, + trust_remote_code=True) + + @classmethod + def download_model(cls, + model_id_or_path: Optional[str] = None, + revision: Optional[str] = None, + ignore_patterns: Optional[List[str]] = None, + **kwargs): + if revision is None or revision == 'master': + revision = 'main' + logger.info(f'Downloading the model from HuggingFace Hub, model_id: {model_id_or_path}') + use_hf_transfer = strtobool(os.environ.get('USE_HF_TRANSFER', 'False')) + if use_hf_transfer: + from huggingface_hub import _snapshot_download + _snapshot_download.HF_HUB_ENABLE_HF_TRANSFER = True + from huggingface_hub import snapshot_download + return snapshot_download( + model_id_or_path, repo_type='model', revision=revision, ignore_patterns=ignore_patterns, **kwargs) + + +def get_hub(use_hf: Optional[bool] = None): + if use_hf is None: + use_hf = True if use_hf_hub() else False + return {True: HFHub, False: MSHub}[use_hf] diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2e1329f534bc3a1d818555fc1b7c4957f9cd71b1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/__init__.py @@ -0,0 +1,39 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING + +from swift.utils.import_utils import _LazyModule + +if TYPE_CHECKING: + from .base import BaseInferEngine + from .grpo_vllm_engine import GRPOVllmEngine + from .infer_client import InferClient + from .infer_engine import InferEngine + from .lmdeploy_engine import LmdeployEngine + from .protocol import ChatCompletionResponse, Function, InferRequest, RequestConfig + from .sglang_engine import SglangEngine + from .transformers_engine import TransformersEngine + from .utils import AdapterRequest, patch_vllm_memory_leak, prepare_generation_config + from .vllm_engine import VllmEngine +else: + _import_structure = { + 'vllm_engine': ['VllmEngine'], + 'grpo_vllm_engine': ['GRPOVllmEngine'], + 'lmdeploy_engine': ['LmdeployEngine'], + 'sglang_engine': ['SglangEngine'], + 'transformers_engine': ['TransformersEngine'], + 'infer_client': ['InferClient'], + 'infer_engine': ['InferEngine'], + 'base': ['BaseInferEngine'], + 'utils': ['prepare_generation_config', 'AdapterRequest', 'patch_vllm_memory_leak'], + 'protocol': ['InferRequest', 'RequestConfig', 'Function', 'ChatCompletionResponse'], + } + + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/base.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/base.py new file mode 100644 index 0000000000000000000000000000000000000000..2e0505871139e77701824782ab32dc025a2ccb4a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/base.py @@ -0,0 +1,58 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from abc import ABC, abstractmethod +from typing import AsyncIterator, Iterator, List, Optional, Union + +from swift.metrics import Metric +from .protocol import ChatCompletionResponse, ChatCompletionStreamResponse, InferRequest, RequestConfig + + +class BaseInferEngine(ABC): + + @abstractmethod + def infer(self, + infer_requests: List[InferRequest], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + use_tqdm: Optional[bool] = None, + **kwargs) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + """ + This method performs inference on a list of inference requests. + + The method takes a list of inference requests and processes them according to the provided configuration. + It can optionally use tqdm for progress visualization and accept additional keyword arguments. + + Args: + infer_requests (List[InferRequest]): A list of inference requests to be processed. + request_config (Optional[RequestConfig]): Configuration for the request, if any. + metrics (Optional[List[Metric]]): A list of usage information to return. + use_tqdm (Optional[bool]): Whether to use tqdm for progress visualization. + **kwargs: Additional keyword arguments. + + Returns: + List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + The result of the inference. + """ + pass + + @abstractmethod + async def infer_async(self, + infer_request: InferRequest, + request_config: Optional[RequestConfig] = None, + **kwargs) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]: + """ + This method performs asynchronous inference on a single inference request. + + The method takes an inference request and processes it according to the provided configuration. + It can accept additional keyword arguments. + + Args: + infer_request (InferRequest): An inference request to be processed. + request_config (Optional[RequestConfig]): Configuration for the request, if any. + **kwargs: Additional keyword arguments. + + Returns: + Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]: The result of + the asynchronous inference. + """ + pass diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/grpo_vllm_engine.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/grpo_vllm_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..d550cefb822773c81f0f561d214a6e83738e848b --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/grpo_vllm_engine.py @@ -0,0 +1,179 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import inspect +import os +import torch +from PIL import Image +from tqdm.asyncio import tqdm_asyncio +from typing import Any, Dict, List, Optional, Union + +from swift.metrics import Metric +from swift.rlhf_trainers.utils import VLLM_LORA_INT_ID, VLLM_LORA_NAME, VLLM_LORA_PATH +from swift.template import Template +from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatMessage, InferRequest, RequestConfig, + RolloutOutput) +from .utils import AdapterRequest +from .vllm_engine import VllmEngine + +try: + os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn' + os.environ['VLLM_ENGINE_ITERATION_TIMEOUT_S'] = '86400' + from vllm.lora.request import LoRARequest +except Exception: + raise + + +class GRPOVllmEngine(VllmEngine): + + def _get_grpo_lora_request(self, adapter_request: Optional[AdapterRequest] = None): + if not adapter_request and self.enable_lora: + lora_loaded = VLLM_LORA_INT_ID in self.engine.list_loras() + if lora_loaded: + adapter_request = LoRARequest( + lora_name=VLLM_LORA_NAME, lora_int_id=VLLM_LORA_INT_ID, lora_path=VLLM_LORA_PATH) + return adapter_request + + async def _get_grpo_lora_request_async(self, adapter_request: Optional[AdapterRequest] = None): + if not adapter_request and self.enable_lora: + loaded_loras = self.engine.list_loras() + if inspect.isawaitable(loaded_loras): + loaded_loras = await loaded_loras + lora_loaded = VLLM_LORA_INT_ID in loaded_loras + if lora_loaded: + adapter_request = LoRARequest( + lora_name=VLLM_LORA_NAME, lora_int_id=VLLM_LORA_INT_ID, lora_path=VLLM_LORA_PATH) + return adapter_request + + def infer( + self, + infer_requests: List[Union[InferRequest, Dict[str, Any]]], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + use_tqdm: Optional[bool] = None, + adapter_request: Optional[AdapterRequest] = None, + ) -> List[RolloutOutput]: + adapter_request = self._get_grpo_lora_request(adapter_request) + + res = super().infer( + infer_requests, + request_config, + metrics, + use_tqdm=use_tqdm, + adapter_request=adapter_request, + ) + if not isinstance(res, list): + res = [res] + for i, result in enumerate(res): + if not isinstance(result, RolloutOutput): + if not isinstance(result, ChatCompletionResponse): + raise TypeError('Result must be a ChatCompletionResponse or RolloutOutput instance.') + res[i] = RolloutOutput(response=result) + + return res + + async def async_infer(self, + infer_requests: List[InferRequest], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + use_tqdm: Optional[bool] = None, + **kwargs) -> List[RolloutOutput]: + if request_config is None: + request_config = RequestConfig() + assert request_config.n == 1 + + tasks = [self.infer_async(infer_request, request_config, **kwargs) for infer_request in infer_requests] + if use_tqdm is None: + use_tqdm = len(infer_requests) > 1 + res = await self._batch_infer_stream(tasks, request_config.stream, use_tqdm, metrics) + + for i, result in enumerate(res): + if not isinstance(result, RolloutOutput): + if not isinstance(result, ChatCompletionResponse): + raise TypeError('Result must be a ChatCompletionResponse or RolloutOutput instance.') + res[i] = RolloutOutput(response=result) + + return res + + async def infer_async( + self, + infer_request: InferRequest, + request_config: Optional[RequestConfig] = None, + *, + adapter_request: Optional[AdapterRequest] = None, + pre_infer_hook=None, + ): + adapter_request = await self._get_grpo_lora_request_async(adapter_request) + return await super().infer_async( + infer_request, + request_config, + adapter_request=adapter_request, + pre_infer_hook=pre_infer_hook, + ) + + async def _batch_infer_stream(self, + tasks, + stream: bool = True, + use_tqdm: bool = True, + metrics: Optional[List[Metric]] = None): + assert not stream + prog_bar = tqdm_asyncio(total=len(tasks), dynamic_ncols=True, disable=not use_tqdm) + + async def _new_run(task): + try: + res = await task + except Exception as e: + if getattr(self, 'strict', True): + raise + res = e + prog_bar.update() + self._update_metrics(res, metrics) + return res + + new_tasks = [_new_run(task) for task in tasks] + return await self.batch_run(new_tasks) + + def _create_chat_completion_response(self, result, inputs, request_config, request_id) -> ChatCompletionResponse: + assert result is not None + num_generated_tokens = sum(len(output.token_ids) for output in result.outputs) + usage_info = self._get_usage_info(len(result.prompt_token_ids), num_generated_tokens) + choices = [] + for output in result.outputs: + output.token_ids = list(output.token_ids) + response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs']) + logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs) + toolcall = self._get_toolcall(response) + + token_ids = output.token_ids if request_config.return_details else None + choice = ChatCompletionResponseChoice( + index=output.index, + message=ChatMessage(role='assistant', content=response, tool_calls=toolcall), + finish_reason=output.finish_reason, + logprobs=logprobs, + token_ids=token_ids, + routed_experts=getattr(output, 'routed_experts', None)) + choices.append(choice) + prompt_token_ids = None + images_size = None + if request_config.return_details: + prompt_token_ids = result.prompt_token_ids + images = inputs['template_inputs'].images + if all(isinstance(image, Image.Image) for image in images): + images_size = [image.size for image in images] + return ChatCompletionResponse( + model=self.model_name, + choices=choices, + usage=usage_info, + id=request_id, + prompt_token_ids=prompt_token_ids, + images_size=images_size) + + def _add_adapter(self, adapter_request: Optional[Union[AdapterRequest, LoRARequest]] = None): + assert self.enable_lora, f'adapter_request: {adapter_request}, self.enable_lora: {self.enable_lora}' + from vllm.lora.request import LoRARequest + if isinstance(adapter_request, AdapterRequest): + return super()._add_adapter(adapter_request) + elif isinstance(adapter_request, LoRARequest): + return adapter_request + else: + raise ValueError(f'Invalid adapter request: {adapter_request}') diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/infer_client.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/infer_client.py new file mode 100644 index 0000000000000000000000000000000000000000..055ecac0cc07ce902e92f9518f6c762838b4426c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/infer_client.py @@ -0,0 +1,158 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import aiohttp +import json +from copy import deepcopy +from dacite import from_dict +from dataclasses import asdict +from requests.exceptions import HTTPError +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union + +from swift.metrics import Metric +from .infer_engine import InferEngine +from .protocol import (ChatCompletionRequest, ChatCompletionResponse, ChatCompletionStreamResponse, InferRequest, + ModelList, RequestConfig) + + +class InferClient(InferEngine): + + def __init__(self, + host: str = '127.0.0.1', + port: int = 8000, + api_key: str = 'EMPTY', + *, + base_url: Optional[str] = None, + timeout: Optional[int] = 86400) -> None: + """ + Initialize the InferClient. + + Args: + host (str): The hostname of the inference server. Defaults to '127.0.0.1'. + port (str): The port of the inference server. Defaults to '8000'. + api_key (str): The API key for authentication. Defaults to 'EMPTY'. + timeout (Optional[int]): The timeout for requests in seconds. Defaults to None. + """ + self.api_key = api_key + self.host = host + self.port = port + self.timeout = timeout + if base_url is None: + base_url = f'http://{self.host}:{self.port}/v1' + self.base_url = base_url + self._models = None + + @property + def models(self): + if self._models is None: + models = [] + for model in self.get_model_list().data: + models.append(model.id) + assert len(models) > 0, f'models: {models}' + self._models = models + return self._models + + def get_model_list(self) -> ModelList: + """Get model list from the inference server. + """ + coro = self.get_model_list_async() + return self.safe_asyncio_run(coro) + + def _get_request_kwargs(self) -> Dict[str, Any]: + request_kwargs = {} + if isinstance(self.timeout, (int, float)) and self.timeout > 0: + request_kwargs['timeout'] = self.timeout + if self.api_key is not None: + request_kwargs['headers'] = {'Authorization': f'Bearer {self.api_key}'} + return request_kwargs + + async def get_model_list_async(self) -> ModelList: + url = f"{self.base_url.rstrip('/')}/models" + async with aiohttp.ClientSession() as session: + async with session.get(url, **self._get_request_kwargs()) as resp: + resp_obj = await resp.json() + return from_dict(ModelList, resp_obj) + + def infer( + self, + infer_requests: List[InferRequest], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + model: Optional[str] = None, + use_tqdm: Optional[bool] = None + ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + """ + Perform inference using the specified model. + + Args: + infer_requests (List[InferRequest]): A list of inference requests. + request_config (Optional[RequestConfig]): Configuration for the request. Defaults to None. + metrics (Optional[List[Metric]]): The usage information to return. Defaults to None. + model (Optional[str]): The model name to be used for inference. Defaults to None. + use_tqdm (Optional[bool]): Whether to use tqdm for progress tracking. Defaults to None. + + Returns: + List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + The inference responses or an iterator of streaming responses. + """ + return super().infer(infer_requests, request_config, metrics, model=model, use_tqdm=use_tqdm) + + @staticmethod + def _prepare_request_data(model: str, infer_request: InferRequest, request_config: RequestConfig) -> Dict[str, Any]: + if not isinstance(infer_request, dict): + infer_request = asdict(infer_request) + res = asdict(ChatCompletionRequest(model, **infer_request, **asdict(request_config))) + # ignore empty + empty_request = ChatCompletionRequest('', []) + for k in list(res.keys()): + if res[k] == getattr(empty_request, k): + res.pop(k) + return res + + @staticmethod + def _parse_stream_data(data: bytes) -> Optional[str]: + data = data.decode(encoding='utf-8') + data = data.strip() + if len(data) == 0: + return + assert data.startswith('data:'), f'data: {data}' + return data[5:].strip() + + async def infer_async( + self, + infer_request: InferRequest, + request_config: Optional[RequestConfig] = None, + *, + model: Optional[str] = None, + ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]: + request_config = deepcopy(request_config or RequestConfig()) + if model is None: + if len(self.models) == 1: + model = self.models[0] + else: + raise ValueError(f'Please explicitly specify the model. Available models: {self.models}.') + url = f"{self.base_url.rstrip('/')}/chat/completions" + + request_data = self._prepare_request_data(model, infer_request, request_config) + if request_config.stream: + + async def _gen_stream() -> AsyncIterator[ChatCompletionStreamResponse]: + async with aiohttp.ClientSession() as session: + async with session.post(url, json=request_data, **self._get_request_kwargs()) as resp: + async for data in resp.content: + data = self._parse_stream_data(data) + if data == '[DONE]': + break + if data is not None: + resp_obj = json.loads(data) + if resp_obj['object'] == 'error': + raise HTTPError(resp_obj['message']) + yield from_dict(ChatCompletionStreamResponse, resp_obj) + + return _gen_stream() + else: + async with aiohttp.ClientSession() as session: + async with session.post(url, json=request_data, **self._get_request_kwargs()) as resp: + resp_obj = await resp.json() + if resp_obj['object'] == 'error': + raise HTTPError(resp_obj['message']) + return from_dict(ChatCompletionResponse, resp_obj) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/infer_engine.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/infer_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..77dd8f6cefca202a41326f957ec3c1cc54cc6e94 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/infer_engine.py @@ -0,0 +1,319 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio +import concurrent.futures +import os +from queue import Queue +from threading import Thread +from tqdm import tqdm +from typing import Any, Dict, Iterator, List, Optional, Union + +from swift.metrics import Metric +from swift.model import get_ckpt_dir +from swift.template import Template, get_template +from swift.utils import Processor, ProcessorMixin, get_logger +from .base import BaseInferEngine +from .protocol import (ChatCompletionMessageToolCall, ChatCompletionResponse, ChatCompletionStreamResponse, + InferRequest, RequestConfig, UsageInfo) + +logger = get_logger() + + +class InferEngine(BaseInferEngine, ProcessorMixin): + + def __init__(self, template: Template): + processor = template.processor + self.template = template + self.template_type = template.template_meta.template_type + self.processor = processor + self.model_info = processor.model_info + self.model_meta = processor.model_meta + self.model_dir = self.model_info.model_dir + self.model_name = self.model_info.model_name + self.max_model_len = self.model_info.max_model_len + self.task_type = self.model_info.task_type + self.config = self.model_info.config + self.max_tokens_offset = 0 + + def _get_template(self, processor: Processor, template_type: Optional[str] = None): + ckpt_dir = get_ckpt_dir(processor.model_info.model_dir, getattr(self, 'adapters', None)) + logger.info('Create the template for the infer_engine') + if ckpt_dir: + from swift.arguments import BaseArguments + args = BaseArguments.from_pretrained(ckpt_dir) + template = args.get_template(processor) + else: + template = get_template(processor, template_type=template_type) + return template + + def _get_stop_words(self, stop_words: List[Union[str, List[int], None]]) -> List[str]: + stop: List[str] = [] + for stop_word in stop_words: + if stop_word is None: + continue + elif isinstance(stop_word, list): + stop_word = self.tokenizer.decode(stop_word) + assert isinstance(stop_word, str) + if stop_word not in stop: + stop.append(stop_word) + return stop + + def _get_stop_token_ids(self, stop_words: List[Union[str, List[int], None]]) -> List[int]: + stop_token_ids: List[int] = [] + for stop_word in stop_words: + if stop_word is None: + continue + if isinstance(stop_word, str): + stop_word = self.tokenizer.encode(stop_word, add_special_tokens=False) + if isinstance(stop_word, list): + if len(stop_word) != 1: + continue + else: + stop_token = stop_word[0] + elif isinstance(stop_word, int): + stop_token = stop_word + assert isinstance(stop_token, int) + if stop_token not in stop_token_ids: + stop_token_ids.append(stop_token) + return stop_token_ids + + def async_iter_to_iter(self, async_iter, prog_bar, metrics) -> Iterator: + queue = Queue() + + async def _run_async_iter(): + try: + async for item in await async_iter: + queue.put(item) + except Exception as e: + queue.put(e) + else: + queue.put(None) + + try: + loop = asyncio.get_event_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + thread = Thread(target=lambda: loop.run_until_complete(_run_async_iter())) + thread.start() + pre_output = None + while True: + output = queue.get() + if output is None or isinstance(output, Exception): + prog_bar.update() + self._update_metrics(pre_output, metrics) + if isinstance(output, Exception) and getattr(self, 'strict', True): + raise output + return + pre_output = output + yield output + + @staticmethod + async def batch_run(tasks): + return await asyncio.gather(*tasks) + + def _batch_infer_stream( + self, + tasks, + stream: bool = True, + use_tqdm: bool = True, + metrics: Optional[List[Metric]] = None + ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + + prog_bar = tqdm(total=len(tasks), dynamic_ncols=True, disable=not use_tqdm) + if stream: + return [self.async_iter_to_iter(task, prog_bar, metrics) for task in tasks] + else: + + async def _new_run(task): + try: + res = await task + except Exception as e: + if getattr(self, 'strict', True): + raise + res = e + prog_bar.update() + self._update_metrics(res, metrics) + return res + + new_tasks = [_new_run(task) for task in tasks] + try: + loop = asyncio.get_event_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + return loop.run_until_complete(self.batch_run(new_tasks)) + + @staticmethod + def _get_usage_info(num_prompt_tokens: int, num_generated_tokens: int) -> UsageInfo: + return UsageInfo( + prompt_tokens=num_prompt_tokens, + completion_tokens=num_generated_tokens, + total_tokens=num_prompt_tokens + num_generated_tokens, + ) + + @staticmethod + def _update_usage_info(origin_use_info: UsageInfo, num_generated_tokens: int) -> UsageInfo: + return UsageInfo( + prompt_tokens=origin_use_info.prompt_tokens, + completion_tokens=origin_use_info.completion_tokens + num_generated_tokens, + total_tokens=origin_use_info.total_tokens + num_generated_tokens, + ) + + @staticmethod + def _update_metrics(result, metrics: Optional[List[Metric]] = None): + if metrics is None: + return result + result_origin = result + if not isinstance(result, (list, tuple)): + result = [result] + for response in result: + if response is None or isinstance(response, Exception): + continue + for metric in metrics: + metric.update(response) + return result_origin + + def infer(self, + infer_requests: List[InferRequest], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + use_tqdm: Optional[bool] = None, + **kwargs) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + if request_config is None: + request_config = RequestConfig() + tasks = [self.infer_async(infer_request, request_config, **kwargs) for infer_request in infer_requests] + if use_tqdm is None: + use_tqdm = not request_config.stream and len(infer_requests) > 1 + return self._batch_infer_stream(tasks, request_config.stream, use_tqdm, metrics) + + def _get_toolcall(self, response: str) -> Optional[List[ChatCompletionMessageToolCall]]: + try: + functions = self.template.agent_template.get_toolcall(response) + except Exception: + functions = None + if functions: + return [ChatCompletionMessageToolCall(function=function) for function in functions] + + @staticmethod + def _get_num_tokens(inputs: Dict[str, Any]) -> int: + if 'input_ids' in inputs: # 1d or 2d + input_ids = inputs['input_ids'] + if isinstance(input_ids, list): + return len(input_ids) + else: + return input_ids.shape[-1] + elif 'inputs_embeds' in inputs: # 2d or 3d + return inputs['inputs_embeds'].shape[-2] + raise ValueError(f'Unable to retrieve input_ids and inputs_embeds. inputs: {inputs}') + + def set_default_max_tokens(self, request_config: RequestConfig, inputs: Dict[str, Any]) -> None: + max_model_len = self.max_model_len + assert isinstance(inputs, dict) + # The num_tokens takes the maximum value from inputs_list. + num_tokens = self._get_num_tokens(inputs) + max_tokens = request_config.max_tokens + if max_model_len is None: + max_model_len = 8192 + logger.warning( + 'The current model is unable to retrieve `max_model_len`. It is set to the default value of 8192.') + max_max_tokens = max_model_len - num_tokens + self.max_tokens_offset + if max_max_tokens <= 0: + raise ValueError( + f'Input length ({num_tokens}) leaves no room for generation with max_model_len ({max_model_len}) ' + f'and max_tokens_offset ({self.max_tokens_offset}). Please shorten the input or increase max_model_len.' + ) + if max_tokens is None: + request_config.max_tokens = max_max_tokens + elif max_max_tokens < request_config.max_tokens: + logger.warning(f'max_model_len({max_model_len}) - num_tokens({num_tokens}) < max_tokens({max_tokens}). ' + f'Setting max_tokens: {max_model_len - num_tokens}') + request_config.max_tokens = max_max_tokens + + def _get_logprobs(self, + logprobs_list: Optional[List[Dict[int, float]]], + token_ids: List[int], + top_logprobs: Optional[int] = None) -> Optional[Dict[str, Any]]: + if logprobs_list is None or len(token_ids) == 0: + return None + if len(token_ids) > 0: + logprobs_list = logprobs_list[-len(token_ids):] + res = [] + for logprobs, token_id in zip(logprobs_list, token_ids): + token = self.tokenizer.decode(token_id) + _res = {'token': token, 'logprob': logprobs[token_id], 'bytes': list(token.encode('utf8'))} + if top_logprobs is not None: + logprobs = {k: logprobs[k] for k in sorted(logprobs, key=lambda k: -logprobs[k])[:top_logprobs]} + res_top_logprobs = [] + for k, logprob in logprobs.items(): + if logprob == float('-inf'): + continue + token = self.tokenizer.decode(k) + res_top_logprobs.append({'token': token, 'logprob': logprob, 'bytes': list(token.encode('utf8'))}) + _res['top_logprobs'] = res_top_logprobs + res.append(_res) + return {'content': res} + + @staticmethod + def _get_finish_reason(max_tokens: int, completion_tokens: int, is_finished: bool): + if is_finished: + if completion_tokens >= max_tokens: + finish_reason = 'length' + else: + finish_reason = 'stop' + else: + finish_reason = None + return finish_reason + + @staticmethod + def thread_run(target, args=(), kwargs=None): + kwargs = kwargs or {} + + def func(target, queue, args, kwargs): + try: + queue.put(target(*args, **kwargs)) + except Exception as e: + queue.put(e) + + queue = Queue() + thread = Thread(target=func, args=(target, queue, args, kwargs)) + thread.start() + thread.join() + result = queue.get() + if isinstance(result, Exception): + raise result + return result + + @staticmethod + def safe_asyncio_run(coro): + + def asyncio_run(core): + return asyncio.run(core) + + return InferEngine.thread_run(asyncio_run, args=(coro, )) + + def _batch_encode(self, infer_requests: List[InferRequest], strict: bool): + max_workers = max(min(32, os.cpu_count(), len(infer_requests)), 1) + error_list = [] + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = [ + executor.submit(self.template.encode, infer_request, return_template_inputs=True) + for infer_request in infer_requests + ] + concurrent.futures.wait(futures) + batched_inputs = [] + for i, future in enumerate(futures): + try: + batched_inputs.append(future.result()) + except Exception as e: + if strict: + raise + error_list.append((i, e)) + continue + return batched_inputs, error_list + + @staticmethod + def _add_error_list(outputs, error_list): + for i, error in error_list: + outputs.insert(i, error) + return outputs diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/lmdeploy_engine.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/lmdeploy_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..61a9409e7b40de20827272fbab13a7b9fde327f5 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/lmdeploy_engine.py @@ -0,0 +1,350 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio +import inspect +import lmdeploy +import os +import time +import torch +from contextlib import contextmanager +from copy import deepcopy +from lmdeploy import PytorchEngineConfig, TurbomindEngineConfig, VisionConfig, pipeline +from lmdeploy.api import autoget_backend_config +from lmdeploy.serve import async_engine +from packaging import version +from PIL import Image +from transformers import GenerationConfig +from transformers.utils.versions import require_version +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union + +from swift.metrics import Metric +from swift.model import get_processor +from swift.template import Template +from swift.utils import get_logger, get_seed, safe_snapshot_download +from .infer_engine import InferEngine +from .patch import patch_auto_config, patch_auto_tokenizer +from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, ChatMessage, DeltaMessage, InferRequest, RequestConfig) +from .utils import InferStreamer + +try: + from lmdeploy import EngineGenerationConfig as LmdeployGenerationConfig +except ImportError: + # compat lmdeploy >= 0.6.* + from lmdeploy import GenerationConfig as LmdeployGenerationConfig + +logger = get_logger() + + +class LmdeployEngine(InferEngine): + + def __init__( + self, + model_id_or_path: str, + *, + template: Optional[Template] = None, + torch_dtype: Optional[torch.dtype] = None, + model_type: Optional[str] = None, + template_type: Optional[str] = None, + use_hf: Optional[bool] = None, + hub_token: Optional[str] = None, + revision: Optional[str] = None, + # engine_kwargs + tp: int = 1, + session_len: Optional[int] = None, + cache_max_entry_count: float = 0.8, + quant_policy: int = 0, # e.g. 4, 8 + vision_batch_size: int = 1, # max_batch_size in VisionConfig + engine_kwargs: Optional[Dict[str, Any]] = None, + devices: Optional[List[int]] = None, + ) -> None: + self.model_id_or_path = model_id_or_path + self.torch_dtype = torch_dtype + self.model_type = model_type + self.use_hf = use_hf + self.hub_token = hub_token + self.revision = revision + self.tp = tp + self.session_len = session_len + self.cache_max_entry_count = cache_max_entry_count + self.quant_policy = quant_policy + self.vision_batch_size = vision_batch_size + self.devices = devices + if template is None: + processor = self._get_processor() + template = self._get_template(processor, template_type=template_type) + else: + safe_snapshot_download( + model_id_or_path, + revision=revision, + download_model=True, + use_hf=use_hf, + ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None), + hub_token=hub_token) + super().__init__(template) + + if self.max_model_len is not None: + self.max_model_len -= 1 + self._prepare_engine_kwargs(engine_kwargs) + self.config.torch_dtype = self.torch_dtype = self.torch_dtype or self.model_info.torch_dtype + self._prepare_engine() + self._load_generation_config() + + def _get_processor(self): + return get_processor( + model_id_or_path=self.model_id_or_path, + torch_dtype=self.torch_dtype, + download_model=True, + model_type=self.model_type, + use_hf=self.use_hf, + hub_token=self.hub_token, + revision=self.revision) + + def _prepare_engine_kwargs(self, engine_kwargs): + if engine_kwargs is None: + engine_kwargs = {} + engine_kwargs['tp'] = self.tp + engine_kwargs['session_len'] = self.session_len + engine_kwargs['cache_max_entry_count'] = self.cache_max_entry_count + engine_kwargs['quant_policy'] = self.quant_policy + if 'devices' in inspect.signature(TurbomindEngineConfig).parameters: + engine_kwargs['devices'] = self.devices + backend_config = TurbomindEngineConfig(**engine_kwargs) + backend_config = autoget_backend_config(self.model_dir, backend_config) + self.backend_config = backend_config + logger.info(f'backend_config: {backend_config}') + + pipeline_kwargs = {} + is_multimodal = self.model_meta.is_multimodal + if is_multimodal: + require_version( + 'lmdeploy<0.9', 'LmdeployEngine will no longer maintain inference for ' + 'multimodal models in lmdeploy>=0.9.') + vision_config = VisionConfig(max_batch_size=self.vision_batch_size) + pipeline_kwargs['vision_config'] = vision_config + logger.info(f'vision_config: {vision_config}') + self.pipeline_kwargs = pipeline_kwargs + + @contextmanager + def _patch_pipeline(self): + _old_best_match_model = async_engine.best_match_model + + def _best_match_model(*args, **kwargs) -> Optional[str]: + return self.model_info.model_type + + async_engine.best_match_model = _best_match_model + try: + yield + finally: + async_engine.best_match_model = _old_best_match_model + + def _prepare_engine(self): + with patch_auto_tokenizer(self.tokenizer), patch_auto_config(self.config), self._patch_pipeline(): + engine = pipeline(self.model_dir, backend_config=self.backend_config, **self.pipeline_kwargs) + self.engine = engine + + def _load_generation_config(self): + generation_config_path = os.path.join(self.model_dir, 'generation_config.json') + if os.path.isfile(generation_config_path): + generation_config = GenerationConfig.from_pretrained(self.model_dir) + kwargs = generation_config.to_dict() + max_new_tokens = kwargs.get('max_new_tokens') + if max_new_tokens is None: + kwargs.pop('max_new_tokens', None) + parameters = inspect.signature(LmdeployGenerationConfig).parameters + for k, v in kwargs.copy().items(): + if k not in parameters or v is None: + kwargs.pop(k) + self.generation_config = LmdeployGenerationConfig(**kwargs) + else: + self.generation_config = LmdeployGenerationConfig() + + def _add_stop_words(self, generation_config: LmdeployGenerationConfig, request_config: RequestConfig) -> None: + template_meta = self.template.template_meta + stop_words = (request_config.stop or []) + (self.generation_config.stop_words or []) + template_meta.stop_words + generation_config.stop_words = self._get_stop_token_ids(stop_words) + # compat lmdeploy >= 0.6.* + generation_config.stop_token_ids = generation_config.stop_words + + def _prepare_generation_config(self, request_config: RequestConfig) -> LmdeployGenerationConfig: + kwargs = {'max_new_tokens': request_config.max_tokens} + for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty']: + new_value = getattr(request_config, key) + if new_value is None: + kwargs[key] = getattr(self.generation_config, key) + else: + kwargs[key] = new_value + if request_config.seed is None: + request_config.seed = get_seed() + kwargs['random_seed'] = request_config.seed + if request_config.temperature == 0: + kwargs['temperature'] = 1 # avoid unnecessary process + kwargs['top_k'] = 1 + + if request_config.logprobs: + kwargs['logprobs'] = 1 + if request_config.top_logprobs is not None: + kwargs['logprobs'] = max(1, request_config.top_logprobs) + + res = LmdeployGenerationConfig(**kwargs) + return res + + async def _infer_stream_async( + self, + inputs: Dict[str, Any], + generation_config: LmdeployGenerationConfig, + request_config: RequestConfig, + ) -> AsyncIterator[ChatCompletionStreamResponse]: + session_id = time.time_ns() + kwargs = {'stream_output': True, 'gen_config': generation_config, 'sequence_start': True, 'sequence_end': True} + if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'): + async with self.engine.model_inst(session_id) as inst: + context = self.engine.safe_run(inst, session_id, **inputs, **kwargs) + else: + context = self.engine.safe_run(session_id) + + infer_streamer = InferStreamer(self.template, template_inputs=inputs['template_inputs']) + token_idx = 0 + async with context as gen: + if version.parse(lmdeploy.__version__) < version.parse('0.6.5'): + generator = await self.engine.get_generator(False, session_id) + gen = generator.async_stream_infer(session_id=session_id, **inputs, **kwargs) + is_finished = False + while not is_finished: + try: + output = await gen.__anext__() + except StopAsyncIteration: + is_finished = True + delta_text = infer_streamer.get_printable_text(output.token_ids, is_finished) + if not delta_text and not is_finished: + continue + + logprobs = self._get_logprobs(output.logprobs, output.token_ids[token_idx:], + request_config.top_logprobs) + token_idx = len(output.token_ids) + + usage_info = self._get_usage_info(len(inputs['input_ids']), output.num_token) + toolcall = None + if is_finished: + toolcall = self._get_toolcall( + self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs'])) + finish_reason = self._get_finish_reason(generation_config.max_new_tokens, output.num_token, + output.status.name == 'FINISH') + choices = [ + ChatCompletionResponseStreamChoice( + index=0, + delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall), + finish_reason=finish_reason, + logprobs=logprobs) + ] + yield ChatCompletionStreamResponse(model=self.model_name, choices=choices, usage=usage_info) + + async def _infer_full_async( + self, + inputs: Dict[str, Any], + generation_config: LmdeployGenerationConfig, + request_config: RequestConfig, + ) -> ChatCompletionResponse: + session_id = time.time_ns() + kwargs = {'stream_output': False, 'gen_config': generation_config, 'sequence_start': True, 'sequence_end': True} + if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'): + async with self.engine.model_inst(session_id) as inst: + async with self.engine.safe_run(inst, session_id, **inputs, **kwargs) as gen: + async for output in gen: + pass + if self.engine.backend == 'pytorch': + # manually end pytorch session + await inst.async_end(session_id) + + else: + async with self.engine.safe_run(session_id): + generator = await self.engine.get_generator(False, session_id) + async for output in generator.async_stream_infer(session_id=session_id, **inputs, **kwargs): + pass + + response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs']) + logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs) + + usage_info = self._get_usage_info(len(inputs['input_ids']), output.num_token) + toolcall = self._get_toolcall(response) + finish_reason = self._get_finish_reason(generation_config.max_new_tokens, output.num_token, + output.status.name == 'FINISH') + token_ids = output.token_ids if request_config.return_details else None + choices = [ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role='assistant', content=response, tool_calls=toolcall), + finish_reason=finish_reason, + logprobs=logprobs, + token_ids=token_ids) + ] + prompt_token_ids = None + images_size = None + if request_config.return_details: + prompt_token_ids = inputs['input_ids'] + images = inputs['template_inputs'].images + if all(isinstance(image, Image.Image) for image in images): + images_size = [image.size for image in images] + return ChatCompletionResponse( + model=self.model_name, + choices=choices, + usage=usage_info, + prompt_token_ids=prompt_token_ids, + images_size=images_size) + + async def infer_async(self, + infer_request: InferRequest, + request_config: Optional[RequestConfig] = None, + *, + pre_infer_hook=None, + **kwargs) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]: + request_config = deepcopy(request_config or RequestConfig()) + self.template.set_mode('lmdeploy') + + loop = asyncio.get_running_loop() + with torch.inference_mode(): + inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True) + images = inputs.pop('images', None) + if images: + if version.parse(lmdeploy.__version__) >= version.parse('0.6.5'): + messages = self.engine._convert_prompts(('', images)) + messages = await self.engine.async_convert_to_pil_images(messages) + results = await self.engine.vl_encoder.preprocess(messages) + if self.engine.backend == 'turbomind': + results = await self.engine.vl_encoder.async_infer(results) + inputs['images'] = [result['content'] for result in results if result['role'] == 'forward'][0] + await self.template.prepare_lmdeploy_turbomind_inputs(inputs) + else: + inputs['images'] = results[1]['content'] + await self.template.prepare_lmdeploy_pytorch_inputs(inputs) + else: + inputs['images'] = await self.engine.vl_encoder.async_infer(images) + await self.template.prepare_lmdeploy_turbomind_inputs(inputs) + + self.set_default_max_tokens(request_config, inputs) + generation_config = self._prepare_generation_config(request_config) + self._add_stop_words(generation_config, request_config) + kwargs.update({'inputs': inputs, 'generation_config': generation_config, 'request_config': request_config}) + if pre_infer_hook: + kwargs = pre_infer_hook(kwargs) + if request_config.stream: + return self._infer_stream_async(**kwargs) + else: + return await self._infer_full_async(**kwargs) + + def _batch_infer_stream(self, *args, **kwargs): + if hasattr(self.engine, 'vl_encoder'): + self.engine.vl_encoder._loop_task = None + if hasattr(self.engine, 'free_insts'): + self.engine.free_insts = None + return super()._batch_infer_stream(*args, **kwargs) + + def infer( + self, + infer_requests: List[InferRequest], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + use_tqdm: Optional[bool] = None, + **kwargs, + ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + return super().infer(infer_requests, request_config, metrics, use_tqdm=use_tqdm, **kwargs) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/patch.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/patch.py new file mode 100644 index 0000000000000000000000000000000000000000..778e6c7ea5a2e0465b0f6383193d3904e100be9a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/patch.py @@ -0,0 +1,34 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from contextlib import contextmanager +from functools import wraps +from transformers import AutoConfig, AutoTokenizer, PretrainedConfig, PreTrainedTokenizerBase + + +@contextmanager +def patch_auto_tokenizer(tokenizer: PreTrainedTokenizerBase): + _old_from_pretrained = AutoTokenizer.from_pretrained + + @wraps(_old_from_pretrained) + def _from_pretrained(*args, **kwargs): + return tokenizer + + AutoTokenizer.from_pretrained = _from_pretrained + try: + yield + finally: + AutoTokenizer.from_pretrained = _old_from_pretrained + + +@contextmanager +def patch_auto_config(config: PretrainedConfig): + _old_from_pretrained = AutoConfig.from_pretrained + + @wraps(_old_from_pretrained) + def _from_pretrained(*args, **kwargs): + return (config, {}) if 'return_unused_kwargs' in kwargs else config + + AutoConfig.from_pretrained = _from_pretrained + try: + yield + finally: + AutoConfig.from_pretrained = _old_from_pretrained diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/protocol.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/protocol.py new file mode 100644 index 0000000000000000000000000000000000000000..df28694d0c673f26a55ff1bba4cccbd4d30b5fb1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/protocol.py @@ -0,0 +1,623 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import base64 +import io +import json +import numpy as np +import os +import time +import uuid +from copy import deepcopy +from dataclasses import asdict, dataclass, field, fields +from PIL import Image +from pydantic import AfterValidator, BaseModel, Field, PlainSerializer, field_validator +from typing import Annotated, Any, Dict, List, Literal, Optional, Tuple, Union + +from swift.template import Messages, Tool +from swift.utils import remove_response + + +def serialize_ndarray(value): + if value is None: + return None + if isinstance(value, np.ndarray): + return { + 'data': base64.b64encode(value.tobytes()).decode('ascii'), + 'shape': value.shape, + 'dtype': str(value.dtype), + '__ndarray__': True + } + return value + + +def deserialize_ndarray(value): + if value is None: + return None + if isinstance(value, dict) and value.get('__ndarray__'): + data = base64.b64decode(value['data']) + return np.frombuffer(data, dtype=value['dtype']).reshape(value['shape']) + return value + + +NumpyArray = Annotated[Any, PlainSerializer(serialize_ndarray, return_type=Dict), AfterValidator(deserialize_ndarray)] + + +@dataclass +class InferRequest: + """ + Data structure for inference requests. + + Attributes: + messages (Messages): + The input conversation in messages format. Each message is a dict containing at least + a "role" field (e.g., "user", "assistant", "system") and a "content" field. + Example: + [{ + "role": "user", + "content": [ + { + "type": "image", # can also be audio/video + "image": "", + }, + {"type": "text", "text": "Please describe the picture."}, + ], + }] + The above is equivalent to: + [{"role": "user", "content": "Please describe the picture."}] + with an additional argument: + images = [""] + + images (List[Union[str, Image.Image]]): + Optional, a list of images associated with the request. + Each image can be a URL, local path, base64 string, or PIL.Image object. + + audios (List[str]): + Optional, a list of audio resources associated with the request. + + videos (List[str]): + Optional, a list of video resources associated with the request. + + tools (Optional[List[Tool]]): + An optional list of tools. These should be organized in the agent_template format for + tools requested by the system, for example 'react_en'. + + objects (Dict[str, Any]): + Container for additional multimodal objects, grouped by type (key). + """ + messages: Messages + + images: List[Union[str, Image.Image]] = field(default_factory=list) + audios: List[str] = field(default_factory=list) + videos: List[str] = field(default_factory=list) + + tools: Optional[List[Tool]] = None + objects: Dict[str, Any] = field(default_factory=dict) + chat_template_kwargs: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + for key in ['images', 'audios', 'videos']: + val = getattr(self, key) + if isinstance(val, str): + setattr(self, key, [val]) + assert isinstance(self.messages, list), f'messages: {self.messages}' + + @staticmethod + def remove_response(messages) -> Optional[str]: + return remove_response(messages) + + @staticmethod + def _to_printable(obj, key: Optional[str] = None): + if isinstance(obj, str) and key not in {'content', 'text'} and len(obj) >= 1000: + return f'<<>>' + elif isinstance(obj, list): + res = [] + for item in obj: + res.append(InferRequest._to_printable(item)) + return res + elif isinstance(obj, dict): + res = {} + for k, v in obj.items(): + res[k] = InferRequest._to_printable(v, key=k) + return res + return obj + + def to_printable(self): + return InferRequest._to_printable(asdict(self)) + + +@dataclass +class RolloutInferRequest(InferRequest): + """ + An inference request class for rollout scenarios. + + This class extends `InferRequest` and specifically overrides the `images` attribute + to be a list of strings for compatibility with POST requests. Each string may + represent an image URL or a Base64-encoded image. + + Inherits all fields from `InferRequest`: + messages (Messages): + Input conversation messages, supporting multimodal content. + audios (List[str]): + List of audio resources associated with the request. + videos (List[str]): + List of video resources associated with the request. + tools (Optional[List[Tool]]): + List of tools, organized by the agent template (e.g. 'react_en'). + objects (Dict[str, Any]): + Optional container for additional multimodal objects. + + Additional / Overridden fields: + images (List[str]): + List of image resources, each as a string (URL or base64). + data_dict (Dict): + Optional dictionary for extra request data. + uuid (Optional[str]): + Optional unique identifier for this request instance. + """ + images: List[str] = field(default_factory=list) + data_dict: Dict = field(default_factory=dict) + uuid: Optional[str] = None + + +def random_uuid() -> str: + return str(uuid.uuid4().hex) + + +@dataclass +class Model: + id: str # model_type + + object: str = 'model' + created: int = field(default_factory=lambda: int(time.time())) + owned_by: str = 'ms-swift' + + +@dataclass +class ModelList: + data: List[Model] + object: str = 'list' + + +@dataclass +class RequestConfig: + """NOTE: The following behavior is inconsistent with the OpenAI API. + Default values for OpenAI: + temperature = 1. + top_k = -1 + top_p = 1. + repetition_penalty = 1. + """ + max_tokens: Optional[int] = None # None: max_model_len - num_tokens + # None: use deploy_args + temperature: Optional[float] = None + top_k: Optional[int] = None + top_p: Optional[float] = None + min_p: Optional[float] = None + repetition_penalty: Optional[float] = None + num_beams: int = 1 + stop: Optional[List[str]] = field(default_factory=list) + + seed: Optional[int] = None + stream: bool = False + logprobs: bool = False + top_logprobs: Optional[int] = None + prompt_logprobs: Optional[int] = None + + n: int = 1 + best_of: Optional[int] = None + presence_penalty: float = 0. + frequency_penalty: float = 0. + length_penalty: float = 1. + # Return token_ids additionally (non-stream) + return_details: bool = False + # vLLM structured outputs (guided decoding) + structured_outputs_regex: Optional[str] = None + + def __post_init__(self): + if self.stop is None: + self.stop = [] + + +@dataclass +class CompletionRequestMixin: + model: str + prompt: str + + +@dataclass +class EmbeddingRequestMixin: + input: str + model: str + encoding_format: Literal['float', 'base64'] = 'float' + + +@dataclass +class ChatCompletionRequestMixin: + model: str + messages: Messages + tools: Optional[List[Tool]] = None + tool_choice: Optional[Union[str, Dict]] = None + chat_template_kwargs: Dict[str, Any] = field(default_factory=dict) + + def __post_init__(self): + if self.tool_choice is None: + self.tool_choice = 'none' if self.tools is None else 'auto' + + if self.tools: + if self.tool_choice == 'none': + self.tools = None + elif isinstance(self.tool_choice, dict): + name = self.tool_choice['function']['name'] + tool = next(tool for tool in self.tools if tool['function']['name'] == name) + if tool is None: + raise ValueError(f"Tool choice '{name}' not found in tools.") + self.tools = [tool] + + +@dataclass +class MultiModalRequestMixin: + images: List[str] = field(default_factory=list) + audios: List[str] = field(default_factory=list) + videos: List[str] = field(default_factory=list) + objects: Dict[str, Any] = field(default_factory=dict) + + @staticmethod + def to_base64(mm_data: Union[str, Image.Image, bytes]) -> str: + if isinstance(mm_data, dict) and 'bytes' in mm_data: + mm_data = mm_data['bytes'] or mm_data['path'] + if isinstance(mm_data, str) and not os.path.isfile(mm_data): + # base64 or url + return mm_data + if isinstance(mm_data, str): + # local_path + with open(mm_data, 'rb') as f: + bytes_ = f.read() + elif isinstance(mm_data, Image.Image): + bytes_io = io.BytesIO() + mm_data.save(bytes_io, format='png') + bytes_ = bytes_io.getvalue() + else: + bytes_ = mm_data + img_base64: str = base64.b64encode(bytes_).decode('utf-8') + return img_base64 + + def __post_init__(self): + for key in ['images', 'audios', 'videos']: + values = getattr(self, key) + if isinstance(values, str): + values = [values] + setattr(self, key, values) + for i, val in enumerate(values): + values[i] = self.to_base64(val) + + +@dataclass +class CompletionRequest(RequestConfig, MultiModalRequestMixin, CompletionRequestMixin): + + def __post_init__(self): + RequestConfig.__post_init__(self) + MultiModalRequestMixin.__post_init__(self) + + +@dataclass +class EmbeddingRequest(RequestConfig, MultiModalRequestMixin, EmbeddingRequestMixin): + + def __post_init__(self): + RequestConfig.__post_init__(self) + MultiModalRequestMixin.__post_init__(self) + + def parse(self) -> Tuple['InferRequest', 'RequestConfig']: + data = asdict(self) + res = [] + for cls_type in [InferRequest, RequestConfig]: + parameters = set(f.name for f in fields(cls_type)) + _data = {k: v for k, v in data.items() if k in parameters} + res.append(cls_type(**_data)) + return tuple(res) + + +@dataclass +class ChatCompletionRequest(RequestConfig, MultiModalRequestMixin, ChatCompletionRequestMixin): + + def __post_init__(self): + RequestConfig.__post_init__(self) + MultiModalRequestMixin.__post_init__(self) + ChatCompletionRequestMixin.__post_init__(self) + self.convert_to_base64() + + def convert_to_base64(self): + for message in self.messages: + content = message['content'] + if isinstance(content, str): + continue + for item in content: + key: str = item['type'] + if key == 'text': + continue + + key_origin = key + value = item[key] + if key.endswith('_url'): + key = key[:-len('_url')] + is_dict = False + if isinstance(value, dict): + is_dict = True + value = value['url'] + if isinstance(value, str) and (value.startswith('data:') or value.startswith('http') + or len(value) > 200): + continue + + # local_path / PIL.Image + if isinstance(value, str) and os.path.isfile(value): + suffix = os.path.splitext(value)[1][1:].lower() + elif isinstance(value, Image.Image): + suffix = 'jpeg' + else: + raise ValueError(f'value: {value}') + mm_data_base64 = self.to_base64(value) + new_value = f'data:{key}/{suffix};base64,{mm_data_base64}' + if is_dict: + new_value = {'url': new_value} + item[key_origin] = new_value + + def parse(self) -> Tuple['InferRequest', 'RequestConfig']: + data = asdict(self) + res = [] + for cls_type in [InferRequest, RequestConfig]: + parameters = set(f.name for f in fields(cls_type)) + _data = {k: v for k, v in data.items() if k in parameters} + res.append(cls_type(**_data)) + return tuple(res) + + @classmethod + def from_cmpl_request(cls, cmpl_request: Union[CompletionRequest, EmbeddingRequest]) -> 'ChatCompletionRequest': + cmpl_request = asdict(cmpl_request) + if 'prompt' in cmpl_request: + prompt = cmpl_request.pop('prompt') + else: + prompt = cmpl_request.pop('input') + cmpl_request['messages'] = [{'role': 'user', 'content': prompt}] + if 'encoding_format' in cmpl_request: + cmpl_request.pop('encoding_format') + return cls(**cmpl_request) + + +@dataclass +class UsageInfo: + prompt_tokens: int + completion_tokens: int + total_tokens: int + + +@dataclass +class Function: + name: str + arguments: Optional[Union[str, Any]] + + def __post_init__(self): + if not isinstance(self.arguments, str): + self.arguments = json.dumps(self.arguments, ensure_ascii=False) + self.name = self.name.strip() + self.arguments = self.arguments.strip() + + +@dataclass +class ChatCompletionMessageToolCall: + function: Function + type: str = 'function' + id: str = field(default_factory=lambda: f'toolcall-{random_uuid()}') + + +@dataclass +class ChatMessage: + role: Literal['system', 'user', 'assistant'] + content: Union[str, List[Dict[str, Any]], int, float, List[float]] + tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None + reasoning_content: Optional[str] = None + + +@dataclass +class ChatCompletionResponseChoice: + index: int + message: ChatMessage + finish_reason: Literal['stop', 'length', None] + logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None + token_ids: Optional[List[int]] = None + routed_experts: Optional[NumpyArray] = None + + def to_cmpl_choice(self) -> 'CompletionResponseChoice': + self = deepcopy(self) + assert not self.message.tool_calls, f'message: {self.message}' + return CompletionResponseChoice(self.index, self.message.content, self.finish_reason, self.logprobs) + + +@dataclass +class EmbeddingResponseData: + object: str = 'embedding' + index: int = 0 + embedding: List[str] = field(default_factory=lambda: []) + + +@dataclass +class EmbeddingResponse: + model: str + data: List[EmbeddingResponseData] + usage: UsageInfo + id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}') + object: str = 'list' + created: int = field(default_factory=lambda: int(time.time())) + + +@dataclass +class CompletionResponseChoice: + index: int + text: str + finish_reason: Literal['stop', 'length', None] + logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None + + +@dataclass +class ChatCompletionResponse: + model: str + choices: List[ChatCompletionResponseChoice] + usage: UsageInfo + id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}') + object: str = 'chat.completion' + created: int = field(default_factory=lambda: int(time.time())) + prompt_token_ids: Optional[List[int]] = None + prompt_logprobs: Optional[List] = None + images_size: Optional[List[Tuple[int, int]]] = None + + def to_cmpl_response(self) -> 'CompletionResponse': + self = deepcopy(self) + choices = [choice.to_cmpl_choice() for choice in self.choices] + id_ = f'cmpl{self.id[len("chatcmpl"):]}' + return CompletionResponse( + self.model, + choices, + self.usage, + id_, + created=self.created, + prompt_token_ids=self.prompt_token_ids, + prompt_logprobs=self.prompt_logprobs, + ) + + +class RolloutOutput(BaseModel): + """ + Output structure for rollout. + + Attributes: + response (ChatCompletionResponse): + The model's response + + messages (Optional[Messages]): + (Optional) Conversation history for the final rollout; required for multi-turn scenarios. + NOTE: + - If provided, this messages sequence will overwrite the original messages. + - If not provided, 'response' will be appended as the latest turn in the original messages. + - For multi-turn training, you need to manually return the updated messages, including the full history. + - The messages should include the latest assistant response as the final message. + + response_token_ids (Optional[List[List[int]]]): + (Optional) Token IDs generated at each rollout turn. + If provided, the training process will skip tokenizing the response. + + response_loss_mask (Optional[List[List[int]]]): + (Optional) Loss masks corresponding to each rollout turn. + If provided, the training process will skip computing loss masks for the response (as controlled by the `loss_scale` parameter). # noqa + + rollout_infos (Dict[str, Any]): + (Optional) Additional rollout information. This must be JSON-serializable. + """ + response: ChatCompletionResponse + # multi turn + messages: Optional[Messages] = None + response_token_ids: List[List[int]] = Field(default_factory=list) + response_loss_mask: List[List[int]] = Field(default_factory=list) + rollout_infos: Dict[str, Any] = Field(default_factory=dict) + # rollout logprobs for each turn (used for rollout importance sampling correction in multi-turn scenarios) + rollout_logprobs: List[List[float]] = Field(default_factory=list) + + prompt_logprobs: Optional[List] = None + + @field_validator('response_token_ids', 'response_loss_mask', 'rollout_logprobs', mode='before') + @classmethod + def _wrap_flat_list(cls, v): + if isinstance(v, list) and v and isinstance(v[0], (int, float)): + return [v] + return v + + def model_post_init(self, __context): + # Ensure multimodal data in rollout_infos is serializable (e.g., images to base64) + super().model_post_init(__context) + self.mminfo_to_serializable() + + def mminfo_to_serializable(self): + mm_keys = ['images', 'audios', 'videos'] + + for key, values in self.rollout_infos.items(): + if key in mm_keys: + if not isinstance(values, list): + values = [values] + for i, value in enumerate(values): + values[i] = MultiModalRequestMixin.to_base64(value) + self.rollout_infos[key] = values + + +@dataclass +class CompletionResponse: + model: str + choices: List[CompletionResponseChoice] + usage: UsageInfo + id: str = field(default_factory=lambda: f'cmpl-{random_uuid()}') + object: str = 'text_completion' + created: int = field(default_factory=lambda: int(time.time())) + prompt_token_ids: Optional[List[int]] = None + prompt_logprobs: Optional[List] = None + + +@dataclass +class DeltaMessage: + role: Literal['system', 'user', 'assistant', None] = None + content: Optional[str] = None + tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None + reasoning_content: Optional[str] = None + + +@dataclass +class ChatCompletionResponseStreamChoice: + index: int + delta: DeltaMessage + finish_reason: Literal['stop', 'length', None] + logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None + + def to_cmpl_choice(self) -> 'CompletionResponseStreamChoice': + self = deepcopy(self) + assert not self.delta.tool_calls + return CompletionResponseStreamChoice(self.index, self.delta.content, self.finish_reason, self.logprobs) + + +@dataclass +class CompletionResponseStreamChoice: + index: int + text: str + finish_reason: Literal['stop', 'length', None] + logprobs: Optional[Dict[str, List[Dict[str, Any]]]] = None + + +@dataclass +class ChatCompletionStreamResponse: + model: str + choices: List[ChatCompletionResponseStreamChoice] + usage: Optional[UsageInfo] = None + id: str = field(default_factory=lambda: f'chatcmpl-{random_uuid()}') + object: str = 'chat.completion.chunk' + created: int = field(default_factory=lambda: int(time.time())) + + def to_cmpl_response(self) -> 'CompletionStreamResponse': + self = deepcopy(self) + choices = [choice.to_cmpl_choice() for choice in self.choices] + id_ = f'cmpl{self.id[len("chatcmpl"):]}' + return CompletionStreamResponse(self.model, choices, self.usage, id_, created=self.created) + + +@dataclass +class CompletionStreamResponse: + model: str + choices: List[CompletionResponseStreamChoice] + usage: Optional[UsageInfo] = None + id: str = field(default_factory=lambda: f'cmpl-{random_uuid()}') + object: str = 'text_completion.chunk' + created: int = field(default_factory=lambda: int(time.time())) + + +class InitCommunicatorRequest(BaseModel): + host: str + port: int + world_size: int + + +class UpdateWeightsRequest(BaseModel): + name: str + dtype: str + shape: list[int] diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/sglang_engine.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/sglang_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..15045e9881ae1b3361786a7fbbc13c168722cd55 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/sglang_engine.py @@ -0,0 +1,302 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio +import inspect +import os +import sglang as sgl +import torch +from copy import deepcopy +from PIL import Image +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.server_args import ServerArgs +from transformers import GenerationConfig +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union + +from swift.metrics import Metric +from swift.model import get_processor +from swift.template import Template +from swift.utils import get_logger, safe_snapshot_download +from .infer_engine import InferEngine +from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, ChatMessage, DeltaMessage, EmbeddingResponse, + EmbeddingResponseData, InferRequest, RequestConfig, random_uuid) +from .utils import InferStreamer + +logger = get_logger() + + +class SglangEngine(InferEngine): + + def __init__( + self, + model_id_or_path: str, + *, + template: Optional[Template] = None, + torch_dtype: Optional[torch.dtype] = None, + model_type: Optional[str] = None, + template_type: Optional[str] = None, + use_hf: Optional[bool] = None, + hub_token: Optional[str] = None, + revision: Optional[str] = None, + # engine kwargs + tp_size: int = 1, + pp_size: int = 1, + dp_size: int = 1, + ep_size: int = 1, + enable_ep_moe: bool = False, + mem_fraction_static: Optional[float] = None, + context_length: Optional[int] = None, + disable_cuda_graph: bool = False, + quantization: Optional[str] = None, + task_type: Optional[str] = None, + kv_cache_dtype: str = 'auto', + enable_dp_attention: bool = False, + disable_custom_all_reduce: bool = True, + speculative_algorithm: Optional[str] = None, + speculative_num_steps: Optional[int] = None, + speculative_eagle_topk: Optional[int] = None, + speculative_num_draft_tokens: Optional[int] = None, + log_level='error', + engine_kwargs: Optional[Dict[str, Any]] = None, + ): + self.model_id_or_path = model_id_or_path + self.torch_dtype = torch_dtype + self.model_type = model_type + self.use_hf = use_hf + self.hub_token = hub_token + self.revision = revision + self.tp_size = tp_size + self.pp_size = pp_size + self.dp_size = dp_size + self.ep_size = ep_size + self.enable_ep_moe = enable_ep_moe + self.mem_fraction_static = mem_fraction_static + self.context_length = context_length + self.disable_cuda_graph = disable_cuda_graph + self.quantization = quantization + self.task_type = task_type + self.kv_cache_dtype = kv_cache_dtype + self.enable_dp_attention = enable_dp_attention + self.disable_custom_all_reduce = disable_custom_all_reduce + self.speculative_algorithm = speculative_algorithm + self.speculative_num_steps = speculative_num_steps + self.speculative_eagle_topk = speculative_eagle_topk + self.speculative_num_draft_tokens = speculative_num_draft_tokens + self.log_level = log_level + if template is None: + processor = self._get_processor() + template = self._get_template(processor, template_type=template_type) + else: + safe_snapshot_download( + model_id_or_path, + revision=revision, + download_model=True, + use_hf=use_hf, + ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None), + hub_token=hub_token) + super().__init__(template) + self._prepare_server_args(engine_kwargs) + self.engine = sgl.Engine(server_args=self.server_args) + self._load_generation_config() + if speculative_num_draft_tokens is not None: + self.max_tokens_offset = -speculative_num_draft_tokens + + def _get_processor(self): + return get_processor( + model_id_or_path=self.model_id_or_path, + torch_dtype=self.torch_dtype, + download_model=True, + model_type=self.model_type, + use_hf=self.use_hf, + hub_token=self.hub_token, + revision=self.revision, + task_type=self.task_type) + + def _prepare_server_args(self, engine_kwargs): + if engine_kwargs is None: + engine_kwargs = {} + if self.context_length is not None: + self.max_model_len = self.context_length + logger.info(f'Setting max_model_len: {self.context_length}') + if self.max_model_len is not None: + self.max_model_len -= 1 + parameters = inspect.signature(ServerArgs).parameters + if 'pp_size' in parameters: + engine_kwargs['pp_size'] = self.pp_size + if 'enable_ep_moe' in parameters: + engine_kwargs['enable_ep_moe'] = self.enable_ep_moe + self.server_args = ServerArgs( + model_path=self.model_dir, + dtype=self.model_info.torch_dtype, + tp_size=self.tp_size, + dp_size=self.dp_size, + ep_size=self.ep_size, + mem_fraction_static=self.mem_fraction_static, + context_length=self.context_length, + disable_cuda_graph=self.disable_cuda_graph, + quantization=self.quantization, + kv_cache_dtype=self.kv_cache_dtype, + enable_dp_attention=self.enable_dp_attention, + disable_custom_all_reduce=self.disable_custom_all_reduce, + speculative_algorithm=self.speculative_algorithm, + speculative_num_steps=self.speculative_num_steps, + speculative_eagle_topk=self.speculative_eagle_topk, + speculative_num_draft_tokens=self.speculative_num_draft_tokens, + log_level=self.log_level, + skip_tokenizer_init=True, + trust_remote_code=True, + **engine_kwargs, + ) + if self.task_type == 'embedding': + self.server_args.is_embedding = True + + def _load_generation_config(self) -> None: + generation_config_path = os.path.join(self.model_dir, 'generation_config.json') + if os.path.isfile(generation_config_path): + generation_config = GenerationConfig.from_pretrained(self.model_dir) + else: + generation_config = GenerationConfig() + kwargs = generation_config.to_dict() + top_k = kwargs.get('top_k') + if top_k == 0: + kwargs['top_k'] = -1 + + parameters = inspect.signature(SamplingParams).parameters + self.generation_config = {k: v for k, v in kwargs.items() if k in parameters and v is not None} + + def _prepare_generation_config(self, request_config: RequestConfig) -> Dict[str, Any]: + kwargs = {'max_new_tokens': request_config.max_tokens} + for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty']: + new_value = getattr(request_config, key) + if new_value is None: + kwargs[key] = self.generation_config.get(key) + else: + kwargs[key] = new_value + for key in ['n', 'frequency_penalty', 'presence_penalty']: + kwargs[key] = getattr(request_config, key) + + return kwargs + + def _add_stop_words(self, generation_config: Dict[str, Any], request_config: RequestConfig) -> None: + template_meta = self.template.template_meta + stop_words = (request_config.stop or []) + (self.generation_config.get('stop') or []) + template_meta.stop_words + generation_config['stop_token_ids'] = self._get_stop_token_ids(stop_words) + + def _create_chat_completion_response(self, output, inputs, return_details: bool = False): + assert output is not None + meta_info = output['meta_info'] + usage_info = self._get_usage_info(meta_info['prompt_tokens'], meta_info['completion_tokens']) + response = self.template.decode_generate_ids(output['output_ids'], template_inputs=inputs['template_inputs']) + toolcall = self._get_toolcall(response) + token_ids = output['output_ids'] if return_details else None + choice = ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role='assistant', content=response, tool_calls=toolcall), + finish_reason=meta_info['finish_reason']['type'], + logprobs=None, + token_ids=token_ids) + prompt_token_ids = None + images_size = None + if return_details: + prompt_token_ids = output.get('prompt_token_ids') + images = inputs['template_inputs'].images + if all(isinstance(image, Image.Image) for image in images): + images_size = [image.size for image in images] + return ChatCompletionResponse( + model=self.model_name, + choices=[choice], + usage=usage_info, + id=random_uuid(), + prompt_token_ids=prompt_token_ids, + images_size=images_size) + + def infer( + self, + infer_requests: List[InferRequest], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + use_tqdm: Optional[bool] = None, + ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + return super().infer(infer_requests, request_config, metrics, use_tqdm=use_tqdm) + + async def infer_async(self, + infer_request: InferRequest, + request_config: Optional[RequestConfig] = None, + *, + pre_infer_hook=None, + **kwargs) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]: + request_config = deepcopy(request_config or RequestConfig()) + self.template.set_mode('sglang') + loop = asyncio.get_running_loop() + with torch.inference_mode(): + inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True) + if self.task_type == 'embedding': + inputs.pop('length', None) + self.set_default_max_tokens(request_config, inputs) + generation_config = self._prepare_generation_config(request_config) + self._add_stop_words(generation_config, request_config) + kwargs.update({'inputs': inputs, 'generation_config': generation_config, 'request_config': request_config}) + if pre_infer_hook: + kwargs = pre_infer_hook(kwargs) + if request_config.stream: + return self._infer_stream_async(**kwargs) + elif self.task_type == 'embedding': + kwargs.pop('generation_config', None) + return await self._infer_embedding_async(**kwargs) + else: + return await self._infer_full_async(**kwargs) + + async def _infer_embedding_async(self, inputs: Dict[str, Any], **kwargs) -> EmbeddingResponse: + from sglang.srt.managers.io_struct import EmbeddingReqInput + obj = EmbeddingReqInput( + input_ids=inputs['input_ids'], image_data=inputs.get('images'), audio_data=inputs.get('audios')) + generator = self.engine.tokenizer_manager.generate_request(obj, None) + output = await generator.__anext__() + usage_info = self._get_usage_info(output['meta_info']['prompt_tokens'], 0) + return EmbeddingResponse( + model=self.model_name, + data=[EmbeddingResponseData(embedding=output['embedding'])], + usage=usage_info, + id=random_uuid()) + + async def _infer_full_async(self, inputs: Dict[str, Any], generation_config: Dict[str, Any], + request_config: RequestConfig) -> ChatCompletionResponse: + engine_inputs = {k: v for k, v in inputs.items() if k != 'template_inputs'} + output = await self.engine.async_generate(**engine_inputs, sampling_params=generation_config) + output['prompt_token_ids'] = inputs['input_ids'] + return self._create_chat_completion_response(output, inputs, request_config.return_details) + + async def _infer_stream_async(self, inputs: Dict[str, Any], generation_config: Dict[str, Any], + **kwargs) -> AsyncIterator[ChatCompletionStreamResponse]: + engine_inputs = {k: v for k, v in inputs.items() if k != 'template_inputs'} + result_generator = await self.engine.async_generate( + **engine_inputs, sampling_params=generation_config, stream=True) + infer_streamer = InferStreamer(self.template, template_inputs=inputs['template_inputs']) + async for output in result_generator: + res = self._create_chat_completion_stream_response(output, infer_streamer) + if res is None: + continue + yield res + + def _create_chat_completion_stream_response(self, output, infer_streamer) -> Optional[ChatCompletionStreamResponse]: + assert output is not None + meta_info = output['meta_info'] + finish_reason = meta_info['finish_reason'] + is_finished = finish_reason is not None + delta_text = infer_streamer.get_printable_text(output['output_ids'], is_finished) + if not delta_text and not is_finished: + return + toolcall = None + if is_finished: + finish_reason = finish_reason['type'] + toolcall = self._get_toolcall( + self.template.decode_generate_ids(output['output_ids'], **infer_streamer.decode_kwargs)) + meta_info = output['meta_info'] + usage_info = self._get_usage_info(meta_info['prompt_tokens'], meta_info['completion_tokens']) + # TODO: logprobs + choice = ChatCompletionResponseStreamChoice( + index=0, + delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall), + finish_reason=finish_reason, + logprobs=None) + return ChatCompletionStreamResponse(model=self.model_name, choices=[choice], usage=usage_info) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/transformers_engine.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/transformers_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..ac51cb46c0bf87a4a826ae68193a40f032f01cb8 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/transformers_engine.py @@ -0,0 +1,604 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio +import hashlib +import inspect +import json +import os +import pickle +import time +import torch +import torch.nn.functional as F +import transformers +from copy import deepcopy +from packaging import version +from PIL import Image +from queue import Queue +from threading import Thread +from torch import nn +from tqdm import tqdm +from transformers import GenerationConfig, LogitsProcessorList +from transformers.utils import is_torch_npu_available +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union + +from swift.metrics import Metric +from swift.model import get_model_processor +from swift.template import Template +from swift.tuners import Swift +from swift.utils import get_last_valid_indices, patch_kernels, safe_snapshot_download, to_device +from .infer_engine import InferEngine +from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, ChatMessage, DeltaMessage, EmbeddingResponse, + EmbeddingResponseData, InferRequest, RequestConfig, random_uuid) +from .utils import AdapterRequest, InferStreamer, LogitsStreamer, TokensIteratorStreamer, prepare_generation_config + +_TRANSFORMERS_GE_5_2 = version.parse(transformers.__version__) >= version.parse('5.2.0') +_kernels_patched = False + + +class _GenerationConfig(GenerationConfig): + + def __repr__(self) -> str: + parameters = inspect.signature(self.to_json_string).parameters + kwargs = {} + if 'ignore_metadata' in parameters: + kwargs['ignore_metadata'] = True + gen_kwargs = json.loads(self.to_json_string(**kwargs)) + gen_kwargs.pop('transformers_version', None) + return f'GenerationConfig({gen_kwargs})' + + +class TransformersEngine(InferEngine): + + def __init__( + self, + model: Union[str, nn.Module], + *, + template: Optional[Template] = None, + adapters: Optional[List[str]] = None, + adapter_names: Optional[List[str]] = None, + max_batch_size: int = 1, # 0/1: no limit + reranker_use_activation: bool = True, + # model kwargs + torch_dtype: Optional[torch.dtype] = None, + model_type: Optional[str] = None, + attn_impl: Optional[str] = None, + experts_impl: Optional[str] = None, + device_map: Optional[Union[str, Dict[str, Any]]] = None, + task_type: Optional[str] = None, + quantization_config=None, + model_kwargs: Optional[Dict[str, Any]] = None, + template_type: Optional[str] = None, + # hub kwargs + use_hf: Optional[bool] = None, + revision: Optional[str] = None, + hub_token: Optional[str] = None, + **kwargs): + if isinstance(adapters, str): + adapters = [adapters] + self.adapters = adapters or [] + self.max_batch_size = max_batch_size + self.reranker_use_activation = reranker_use_activation + + self.torch_dtype = torch_dtype + self.model_type = model_type + self.attn_impl = attn_impl + self.experts_impl = experts_impl + self.device_map = device_map + self.task_type = task_type + self.quantization_config = quantization_config + self.model_kwargs = model_kwargs + + self.use_hf = use_hf + self.revision = revision + self.hub_token = hub_token + global _kernels_patched + if _TRANSFORMERS_GE_5_2 and not _kernels_patched: + if use_hf is not None and 'USE_HF' not in os.environ: + os.environ['USE_HF'] = str(use_hf) + patch_kernels() + _kernels_patched = True + if isinstance(model, str): + self.model, processor = self._get_model_processor(model, **kwargs) + template = self._get_template(processor, template_type=template_type) + elif isinstance(model, nn.Module): + self.model = model + if template is None: + raise ValueError('`template` is required when `model` is a nn.Module') + super().__init__(template) + if isinstance(adapter_names, str): + adapter_names = [adapter_names] + if adapter_names and len(adapter_names) != len(self.adapters): + raise ValueError(f'The length of adapter_names ({len(adapter_names)}) must match the length of ' + f'adapters ({len(self.adapters)})') + for i, adapter in enumerate(self.adapters): + adapter_name = None if adapter_names is None else adapter_names[i] + self._add_adapter( + safe_snapshot_download(adapter, use_hf=self.use_hf, hub_token=self.hub_token), + adapter_name=adapter_name) + self.engine = self.model # dummy + self.generation_config = getattr(self.model, 'generation_config', None) + self._queue = Queue() + self._task_pool = {} + self._adapters_pool = {} + self._task_thread = None + + def _get_model_processor(self, model_id_or_path, **kwargs): + return get_model_processor( + model_id_or_path, + torch_dtype=self.torch_dtype, + model_type=self.model_type, + use_hf=self.use_hf, + hub_token=self.hub_token, + revision=self.revision, + device_map=self.device_map, + quantization_config=self.quantization_config, + attn_impl=self.attn_impl, + experts_impl=self.experts_impl, + task_type=self.task_type, + model_kwargs=self.model_kwargs, + **kwargs) + + def _start_infer_worker(self): + self._task_thread = Thread(target=self._infer_worker, daemon=True) + self._task_thread.start() + + def _fetch_infer_requests(self): + while not self._queue.empty(): + infer_request, kwargs, queue = self._queue.get() + info = hashlib.sha256(pickle.dumps((kwargs['request_config']))).hexdigest() + if info not in self._task_pool: + self._task_pool[info] = kwargs, [] + self._task_pool[info][1].append((infer_request, queue)) + if len(self._task_pool) == 0: + return + key, (kwargs, data) = next(iter(self._task_pool.items())) + max_batch_size = self.max_batch_size + if max_batch_size <= 0: + max_batch_size = len(data) + data, remain_data = data[:max_batch_size], data[max_batch_size:] + if remain_data: + self._task_pool[key] = kwargs, remain_data + else: + self._task_pool.pop(key) + kwargs = kwargs.copy() + kwargs['infer_requests'] = [d[0] for d in data] + queue_list = [d[1] for d in data] + return kwargs, queue_list + + def _infer_worker(self): + while True: + time.sleep(0.01) + item = self._fetch_infer_requests() + if item is not None: + kwargs, queue_list = item + request_config = kwargs['request_config'] + res_list_or_gen = self._infer(**kwargs) + if request_config.stream: + finished = False + while not finished: + try: + res_list = next(res_list_or_gen) + except StopIteration: + finished = True + res_list = [None] * len(queue_list) + for (queue, loop), res in zip(queue_list, res_list): + asyncio.run_coroutine_threadsafe(queue.put(res), loop) + else: + for (queue, loop), res in zip(queue_list, res_list_or_gen): + asyncio.run_coroutine_threadsafe(queue.put(res), loop) + + def _add_adapter(self, adapter_path: str, adapter_name: Optional[str] = None) -> None: + self.model = Swift.from_pretrained(self.model, adapter_path, adapter_name) + + def _prepare_generation_config(self, request_config: RequestConfig) -> _GenerationConfig: + generation_config = prepare_generation_config(self.generation_config, request_config, self.tokenizer) + generation_config.return_dict_in_generate = True + if request_config.logprobs: + generation_config.output_logits = True + generation_config.num_return_sequences = request_config.n + return _GenerationConfig(**generation_config.to_dict()) + + def _add_stop_words(self, generation_config: _GenerationConfig, request_config: RequestConfig) -> None: + template_meta = self.template.template_meta + stop_words = (request_config.stop or []) + template_meta.stop_words + generation_config.stop_words = self._get_stop_words(stop_words) + + @staticmethod + def preprocess_logits(batched_logits: Optional[List[torch.Tensor]], batched_generate_ids: torch.Tensor, + top_logprobs: Optional[int]): + top_logprobs = top_logprobs or 1 + batch_size = batched_generate_ids.shape[0] + if batched_logits is None: + return None + batched_logprobs = [] + for i in range(batch_size): + logprobs_list = [] + generate_ids = batched_generate_ids[i] + for j, logits in enumerate(batched_logits): + token = generate_ids[j].item() + logprobs = torch.log_softmax(logits[i], -1) + tokens = [token] + logprobs.argsort(descending=True, dim=-1)[:top_logprobs].tolist() + logprobs_list.append({token: logprobs[token].item() for token in tokens}) + batched_logprobs.append(logprobs_list) + return batched_logprobs + + @staticmethod + def _update_batched_logprobs(batched_logprobs: List[torch.Tensor], logits_streamer: Optional[LogitsStreamer], + generate_ids: torch.Tensor, top_logprobs: int) -> None: + seq_len = generate_ids.shape[1] - len(batched_logprobs[0]) + if logits_streamer is None or seq_len == 0: + return + + res = [] + for i in range(seq_len): + res.append(logits_streamer.queue.get()) + new_batched_logprobs = TransformersEngine.preprocess_logits(res, generate_ids[:, -seq_len:], top_logprobs) + for logprobs, new_logprobs in zip(batched_logprobs, new_batched_logprobs): + logprobs += new_logprobs + + def _infer_stream(self, inputs: Dict[str, Any], *, generation_config: GenerationConfig, + adapter_request: Optional[AdapterRequest], request_config: RequestConfig, + template_inputs) -> Iterator[List[Optional[ChatCompletionStreamResponse]]]: + if generation_config.num_beams != 1: + error_msg = 'Streaming generation does not support beam search.' + raise ValueError(error_msg) + streamer = TokensIteratorStreamer() + generate_kwargs = { + 'generation_config': generation_config, + 'streamer': streamer, + **inputs, + } + adapter_names = self._get_adapter_names(adapter_request) + if adapter_names is not None: + generate_kwargs['adapter_names'] = adapter_names + num_prompt_tokens = self._get_num_tokens(inputs) + + logits_streamer = None + if generation_config.output_logits: + logits_streamer = LogitsStreamer() + generate_kwargs['logits_processor'] = LogitsProcessorList([logits_streamer]) + + def _model_generate(**kwargs): + if is_torch_npu_available(): + torch.npu.set_device(self.model.device) + self.template.generate(self.model, **kwargs) + + generate_kwargs = self.template.prepare_generate_kwargs(generate_kwargs, model=self.model) + thread = Thread(target=_model_generate, kwargs=generate_kwargs) + thread.start() + batch_size = inputs['attention_mask'].shape[0] + all_is_finished = False + is_finished = [False] * batch_size + infer_streamers = [InferStreamer(self.template, template_inputs=template_inputs[i]) for i in range(batch_size)] + request_id_list = [f'chatcmpl-{random_uuid()}' for _ in range(batch_size)] + token_idxs = [0] * batch_size + + raw_batched_generate_ids = None # or torch.Tensor: [batch_size, seq_len] + batched_logprobs = [[] for _ in range(batch_size)] + while not all_is_finished: + try: + batched_tokens = next(streamer) + if batched_tokens.ndim == 1: + batched_tokens = batched_tokens[:, None] + + raw_batched_generate_ids = torch.concat( + [batched_tokens] + if raw_batched_generate_ids is None else [raw_batched_generate_ids, batched_tokens], + dim=1) + except StopIteration: + all_is_finished = True + + batched_generate_ids = self.template.get_generate_ids(raw_batched_generate_ids, num_prompt_tokens) + self._update_batched_logprobs(batched_logprobs, logits_streamer, batched_generate_ids, + request_config.top_logprobs) + + res = [] + for i in range(batched_generate_ids.shape[0]): + if is_finished[i]: + res.append(None) + continue + generate_ids = batched_generate_ids[i] + + # ignore pad_token + masks = generate_ids != self.tokenizer.pad_token_id + generate_ids = generate_ids[masks].tolist() + logprobs_list = None + if batched_logprobs[i]: + logprobs_list = [logprobs for m, logprobs in zip(masks, batched_logprobs[i]) if m.item()] + + is_finished[i] = ( + all_is_finished or is_finished[i] + or len(generate_ids) > 0 and generate_ids[-1] == self.tokenizer.pad_token_id) + delta_text = infer_streamers[i].get_printable_text(generate_ids, is_finished[i]) + if not delta_text and not is_finished[i]: + res.append(None) + continue + logprobs = self._get_logprobs(logprobs_list, generate_ids[token_idxs[i]:], request_config.top_logprobs) + token_idxs[i] = len(generate_ids) + + usage_info = self._get_usage_info(num_prompt_tokens, len(generate_ids)) + toolcall = None + if is_finished[i]: + toolcall = self._get_toolcall( + self.template.decode_generate_ids(generate_ids, template_inputs=template_inputs[i])) + finish_reason = self._get_finish_reason(generation_config.max_new_tokens, usage_info.completion_tokens, + is_finished[i]) + + choices = [ + ChatCompletionResponseStreamChoice( + index=0, + delta=DeltaMessage(role='assistant', content=delta_text, tool_calls=toolcall), + finish_reason=finish_reason, + logprobs=logprobs) + ] + res.append( + ChatCompletionStreamResponse( + model=self.model_name, choices=choices, usage=usage_info, id=request_id_list[i])) + if any(res): + yield res + + def _get_adapter_names(self, adapter_request: Optional[AdapterRequest]) -> Optional[List[str]]: + if adapter_request is None: + if self._adapters_pool: + return ['__base__'] + return + adapter_name = adapter_request.name + if adapter_name not in self._adapters_pool: + self._adapters_pool[adapter_name] = adapter_request + self._add_adapter(adapter_request.path, adapter_name) + return [adapter_name] + + def _infer_forward(self, inputs: Dict[str, Any], adapter_request: Optional[AdapterRequest], + request_config: RequestConfig, **kwargs): + call_kwargs = {} + top_logprobs = request_config.top_logprobs or 20 + adapter_names = self._get_adapter_names(adapter_request) + if adapter_names is not None: + call_kwargs['adapter_names'] = adapter_names + num_prompt_tokens = self._get_num_tokens(inputs) + inputs.pop('labels', None) + output = self.model(**inputs, **call_kwargs) + if hasattr(output, 'logits'): + logits = output.logits + elif 'last_hidden_state' in output: + # embeddings + logits = output['last_hidden_state'] + else: + raise NotImplementedError('Only support `logits` or `hidden_state` in output.') + task_type = self.template.task_type + if task_type == 'seq_cls': + preds, logprobs = self.template.decode_seq_cls(logits, top_logprobs) + elif task_type == 'prm': + preds = self.template.decode_prm(inputs['input_ids'], logits) + logprobs = [None] * len(preds) + elif task_type == 'embedding': + preds = logits + logprobs = [None] * len(preds) + elif task_type in ('reranker', 'generative_reranker'): + if task_type == 'generative_reranker': + attention_mask = inputs.get('attention_mask') + last_valid_indices = -1 if attention_mask is None else get_last_valid_indices(attention_mask) + batch_indices = torch.arange(logits.shape[0], device=logits.device) + logits = logits[batch_indices, last_valid_indices] + preds = logits.float() + if self.reranker_use_activation: + preds = F.sigmoid(preds) + preds = preds.tolist() + logprobs = [None] * len(preds) + else: + raise ValueError(f'Unsupported task_type: {task_type}') + + res = [] + for i, pred in enumerate(preds): + usage_info = self._get_usage_info(num_prompt_tokens, 1) + if task_type == 'embedding': + res.append( + EmbeddingResponse( + model=self.model_name, usage=usage_info, data=[EmbeddingResponseData(embedding=pred.tolist())])) + else: + choices = [ + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role='assistant', content=pred, tool_calls=None), + finish_reason='stop', + logprobs=logprobs[i]) + ] + res.append(ChatCompletionResponse(model=self.model_name, choices=choices, usage=usage_info)) + return res + + def _infer_full(self, inputs: Dict[str, Any], *, generation_config: GenerationConfig, + adapter_request: Optional[AdapterRequest], request_config: RequestConfig, + template_inputs) -> List[ChatCompletionResponse]: + # bos_token TODO: encoder-decoder + generate_kwargs = {'generation_config': generation_config, **inputs} + adapter_names = self._get_adapter_names(adapter_request) + if adapter_names is not None: + generate_kwargs['adapter_names'] = adapter_names + num_prompt_tokens = self._get_num_tokens(inputs) + generate_kwargs = self.template.prepare_generate_kwargs(generate_kwargs, model=self.model) + output = dict(self.template.generate(self.model, **generate_kwargs)) + output.pop('past_key_values', None) + batched_generate_ids = output['sequences'] + batched_generate_ids = self.template.get_generate_ids(batched_generate_ids, num_prompt_tokens) + self.template.debug_logger({'generate_ids': batched_generate_ids}) # debug + batched_logprobs = self.preprocess_logits( + output.get('logits'), batched_generate_ids, request_config.top_logprobs) + + res = [] + num_return_sequences = generation_config.num_return_sequences + for i in range(inputs['attention_mask'].shape[0]): + choices = [] + usage_info = self._get_usage_info(num_prompt_tokens, 0) + for j in range(num_return_sequences): + batched_index = i * num_return_sequences + j + generate_ids = batched_generate_ids[batched_index] + + # ignore pad_token + masks = generate_ids != self.tokenizer.pad_token_id + generate_ids = generate_ids[masks].tolist() + logprobs_list = None + if batched_logprobs is not None: + logprobs_list = [ + logprobs for m, logprobs in zip(masks, batched_logprobs[batched_index]) if m.item() + ] + + logprobs = self._get_logprobs(logprobs_list, generate_ids, request_config.top_logprobs) + usage_info = self._update_usage_info(usage_info, len(generate_ids)) + response = self.template.decode_generate_ids(generate_ids, template_inputs=template_inputs[i]) + finish_reason = self._get_finish_reason(generation_config.max_new_tokens, len(generate_ids), True) + toolcall = self._get_toolcall(response) + token_ids = generate_ids if request_config.return_details else None + choices.append( + ChatCompletionResponseChoice( + index=j, + message=ChatMessage(role='assistant', content=response, tool_calls=toolcall), + finish_reason=finish_reason, + logprobs=logprobs, + token_ids=token_ids)) + prompt_token_ids = None + images_size = None + if request_config.return_details: + if 'input_ids' in inputs: + non_pad_indices = (inputs['input_ids'][i] != self.tokenizer.pad_token_id).nonzero() + if non_pad_indices.numel() > 0: + idx = non_pad_indices.min().item() + prompt_token_ids = inputs['input_ids'][i][idx:].tolist() + if all(isinstance(image, Image.Image) for image in template_inputs[i].images): + images_size = [image.size for image in template_inputs[i].images] + res.append( + ChatCompletionResponse( + model=self.model_name, + choices=choices, + usage=usage_info, + prompt_token_ids=prompt_token_ids, + images_size=images_size)) + return res + + async def infer_async( + self, + infer_request: InferRequest, + request_config: Optional[RequestConfig] = None, + *, + adapter_request: Optional[AdapterRequest] = None, + pre_infer_hook=None, + ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]: + if request_config is None: + request_config = RequestConfig() + queue = asyncio.Queue() + self._queue.put((infer_request, { + 'request_config': request_config, + 'adapter_request': adapter_request, + 'pre_infer_hook': pre_infer_hook + }, (queue, asyncio.get_event_loop()))) + await asyncio.sleep(0) + if self._task_thread is None: + self._start_infer_worker() + if request_config.stream: + + async def _gen_wrapper(): + while True: + item = await queue.get() + await asyncio.sleep(0) + if item is None: + break + yield item + + return _gen_wrapper() + else: + return await queue.get() + + # Ensure `template._post_encode` has no gradient. + @torch.inference_mode() + def _infer( + self, + infer_requests: List[InferRequest], + request_config: RequestConfig, + *, + adapter_request: Optional[AdapterRequest] = None, + pre_infer_hook=None, + ) -> Union[List[ChatCompletionResponse], Iterator[List[Optional[ChatCompletionStreamResponse]]]]: + self.model.eval() + request_config = deepcopy(request_config) + if self.template.use_model: + self.template.model = self.model + + if self.model_info.task_type == 'causal_lm': + self.template.set_mode('transformers') + + batched_inputs, error_list = self._batch_encode(infer_requests, strict=getattr(self, 'strict', True)) + if len(batched_inputs) > 0: + template_inputs = [inputs.pop('template_inputs') for inputs in batched_inputs] + inputs = to_device(self.template.data_collator(batched_inputs), self.model.device) + self.template.debug_logger(inputs) # debug + if self.model_meta.is_multimodal: + _, inputs = self.template.pre_forward_hook(self.model, None, inputs) + if self.model_info.task_type == 'causal_lm': + self.set_default_max_tokens(request_config, inputs) + generation_config = self._prepare_generation_config(request_config) + self._add_stop_words(generation_config, request_config) + else: + generation_config = request_config + + kwargs = { + 'inputs': inputs, + 'generation_config': generation_config, + 'adapter_request': adapter_request, + 'request_config': request_config, + 'template_inputs': template_inputs, + } + if pre_infer_hook: + kwargs = pre_infer_hook(kwargs) + else: + kwargs = {} + if request_config.stream: + + def _gen_wrapper(): + if len(kwargs) > 0: + for res in self._infer_stream(**kwargs): + yield self._add_error_list(res, error_list) + else: + yield self._add_error_list([], error_list) + + return _gen_wrapper() + else: + if len(kwargs) > 0: + infer_func = self._infer_forward if self.template.task_type in { + 'seq_cls', 'prm', 'embedding', 'reranker', 'generative_reranker' + } else self._infer_full + res = infer_func(**kwargs) + else: + res = [] + return self._add_error_list(res, error_list) + + def infer( + self, + infer_requests: List[InferRequest], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + use_tqdm: Optional[bool] = None, + adapter_request: Optional[AdapterRequest] = None + ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + if request_config is None: + request_config = RequestConfig() + if request_config.stream: + return super().infer( + infer_requests, request_config, metrics, use_tqdm=use_tqdm, adapter_request=adapter_request) + # Has higher stability than calling super().infer + if use_tqdm is None: + use_tqdm = not request_config.stream and len(infer_requests) > 1 + prog_bar = tqdm(total=len(infer_requests), dynamic_ncols=True, disable=not use_tqdm) + # If self.max_batch_size <= 0, then process all infer_requests at once. + max_batch_size = self.max_batch_size + if max_batch_size <= 0: + max_batch_size = len(infer_requests) + res = [] + i = 0 + while i < len(infer_requests): + infer_requests_samples = infer_requests[i:i + max_batch_size] + res += self._infer(infer_requests_samples, request_config, adapter_request=adapter_request) + i += max_batch_size + prog_bar.update(len(infer_requests_samples)) + prog_bar.close() + self._update_metrics(res, metrics) + return res diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..34eb5d100983f36d564f5f68c5bade93d59de955 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/utils.py @@ -0,0 +1,633 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +import os +import re +import torch +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from itertools import repeat +from packaging import version +from queue import Queue +from transformers import GenerationConfig, LogitsProcessor +from transformers.generation.streamers import BaseStreamer +from typing import List, Optional, Union + +from swift.model.register import fix_do_sample_warning +from swift.utils import get_device, synchronize +from .protocol import RequestConfig + + +@dataclass +class AdapterRequest: + name: str + path: str + + +class InferTools: + + @staticmethod + def _is_chinese_char(cp: int) -> bool: + """Checks whether CP is the codepoint of a CJK character.""" + # copy from transformers.generation.streamers.TextStreamer + if ((0x4E00 <= cp <= 0x9FFF) or (0x3400 <= cp <= 0x4DBF) or (0x20000 <= cp <= 0x2A6DF) + or (0x2A700 <= cp <= 0x2B73F) or (0x2B740 <= cp <= 0x2B81F) or (0x2B820 <= cp <= 0x2CEAF) + or (0xF900 <= cp <= 0xFAFF) or (0x2F800 <= cp <= 0x2FA1F)): + return True + + return False + + +class InferStreamer(InferTools): + + def __init__(self, template, **decode_kwargs): + self.template = template + self.tokenizer = template.tokenizer + + self.cache_idx = 0 # token idx + self.print_idx = 0 + self.decode_kwargs = decode_kwargs + self.first_num_space = -1 # The number of whitespace characters before the first token. + self.first_token = True + + def _align_blank_suffix(self, response: str) -> str: + # Avoid the occurrence of repeated words in sentence. + cur_num_space = len(response) - len(response.lstrip(' ')) + if self.first_num_space == -1: + self.first_num_space = cur_num_space + elif cur_num_space < self.first_num_space: + response = ' ' * (self.first_num_space - cur_num_space) + response + elif cur_num_space > self.first_num_space: + response = response[cur_num_space - self.first_num_space:] + return response + + def _get_response(self, response: str, is_finished: bool, token_len: int) -> str: + # After the symbol for a new line, we flush the cache. + if self.first_token: + printable_text = response + self.first_token = False + elif response.endswith('\n') or is_finished: + printable_text = response[self.print_idx:] + self.cache_idx += token_len + self.first_num_space = -1 + self.print_idx = 0 + # If the last token is a CJK character, we print the characters. + elif len(response) > 0 and self._is_chinese_char(ord(response[-1])): + printable_text = response[self.print_idx:] + self.print_idx += len(printable_text) + # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words, + # which may change with the subsequent token -- there are probably smarter ways to do this!) + else: + printable_text = response[self.print_idx:response.rfind(' ') + 1] + self.print_idx += len(printable_text) + return printable_text + + def get_printable_text(self, raw_tokens: List[int], is_finished: bool) -> str: + raw_tokens = raw_tokens[self.cache_idx:] + if self.first_token: + raw_tokens = [] + response = self.template.decode_generate_ids( + raw_tokens, is_finished=is_finished, first_token=self.first_token, **self.decode_kwargs) + response = self._align_blank_suffix(response) + return self._get_response(response, is_finished, len(raw_tokens)) + + +class StreamerMixin: + + def __init__(self): + self.queue = Queue() + + def __iter__(self): + return self + + def __next__(self) -> torch.Tensor: + value = self.queue.get() + if value is None: + raise StopIteration() + else: + return value + + +class TokensIteratorStreamer(StreamerMixin, BaseStreamer): + + def put(self, value: torch.Tensor) -> None: + self.queue.put(value) + + def end(self) -> None: + self.queue.put(None) + + +class LogitsStreamer(LogitsProcessor): + + def __init__(self): + self.queue = Queue() + + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + self.queue.put(scores) + return scores + + +def _set_generation_config_default_value(model_generation_config: GenerationConfig, + generation_config: GenerationConfig) -> GenerationConfig: + for k, v in model_generation_config.to_dict().items(): + new_v = getattr(generation_config, k, None) + if k in ['max_length']: + continue + if k in ['no_repeat_ngram_size'] or v is not None and new_v is None: + setattr(generation_config, k, v) + return generation_config + + +def prepare_generation_config(model_generation_config: Optional[GenerationConfig], request_config: RequestConfig, + tokenizer) -> Optional[GenerationConfig]: + if model_generation_config is None or request_config is None: + return model_generation_config + kwargs = {'max_new_tokens': request_config.max_tokens} + # not use: 'n', 'best_of', 'frequency_penalty', 'presence_penalty' + for key in ['length_penalty']: + kwargs[key] = getattr(request_config, key) + for key in ['temperature', 'top_k', 'top_p', 'repetition_penalty', 'num_beams']: + new_value = getattr(request_config, key) + if new_value is None: + kwargs[key] = getattr(model_generation_config, key, None) + else: + kwargs[key] = new_value + + if kwargs.get('top_k') is not None and kwargs['top_k'] <= 0: + kwargs['top_k'] = None + + if not getattr(model_generation_config, 'do_sample', False) and request_config.temperature in {0, None}: + kwargs['temperature'] = 0 + if kwargs['temperature'] == 0: + kwargs['do_sample'] = False + kwargs['temperature'] = 1 + kwargs['top_p'] = 1 + kwargs['top_k'] = 50 + else: + kwargs['do_sample'] = True + generation_config = GenerationConfig(**kwargs) + generation_config = _set_generation_config_default_value(model_generation_config, generation_config) + fix_do_sample_warning(generation_config) + + if generation_config.eos_token_id is None: + generation_config.eos_token_id = tokenizer.eos_token_id + generation_config.pad_token_id = tokenizer.pad_token_id + return generation_config + + +def patch_lmdeploy(load_weights=False): + """This patch allows lmdeploy selects device and reload state_dict""" + import lmdeploy + assert version.parse(lmdeploy.__version__) >= version.parse('0.7.0') + from lmdeploy.messages import TurbomindEngineConfig + from lmdeploy.turbomind.deploy import loader + from lmdeploy.turbomind.deploy.loader import create_loader + from lmdeploy.turbomind.deploy.source_model import llama + + def _create_loader(model_path: str, pattern: str): + if not isinstance(model_path, (str, os.PathLike)): + + def generate(): + generator = OrderedDict() + model_dict = {} + if not isinstance(model_path, dict): + for key, value in list(model_path): + model_dict[key] = value + else: + model_dict = model_path + for key, value in model_dict.items(): + match = re.findall(pattern, key) + if not match: + if -1 not in generator: + generator[-1] = {} + generator[-1][key] = value + else: + layer = int(match[0]) + if layer not in generator: + generator[layer] = {} + generator[layer][key] = value + return generator + + return generate() + else: + return create_loader(model_path, pattern) + + loader.create_loader = _create_loader + llama.create_loader = _create_loader + + TurbomindEngineConfig.devices = [0] + + from lmdeploy.turbomind.turbomind import TurboMind + from lmdeploy.turbomind.utils import ModelSource + + @contextmanager + def patch_threadpool_map(): + ThreadPoolExecutor.map_origin = ThreadPoolExecutor.map + ThreadPoolExecutor.map = lambda *args, **kwargs: [] + yield + ThreadPoolExecutor.map = ThreadPoolExecutor.map_origin + del ThreadPoolExecutor.map_origin + + @contextmanager + def tm_model_context(self): + + def _get_tm_model(model_path, + model_name, + chat_template_name, + engine_config: TurbomindEngineConfig, + group_size: int = None, + out_dir: str = None): + from lmdeploy.turbomind.deploy.converter import get_tm_model_origin + tm_model = get_tm_model_origin(model_path, model_name, chat_template_name, engine_config, group_size, + out_dir) + self.tm_model = tm_model + return tm_model + + from lmdeploy.turbomind.deploy import converter + converter.get_tm_model_origin = converter.get_tm_model + converter.get_tm_model = _get_tm_model + yield + converter.get_tm_model = converter.get_tm_model_origin + del converter.get_tm_model_origin + + def __init__(self, + model_path: str, + tokenizer: object, + model_name: str = None, + chat_template_name: str = None, + engine_config: TurbomindEngineConfig = None, + model_source: ModelSource = ModelSource.WORKSPACE, + **kwargs): + self.gpu_list = engine_config.devices + with patch_threadpool_map(), tm_model_context(self): + self.__origin_init__(model_path, tokenizer, model_name, chat_template_name, engine_config, model_source, + **kwargs) + + with ThreadPoolExecutor(max_workers=self.gpu_count) as e: + ranks = [self.node_id * self.gpu_count + device_id for device_id in range(self.gpu_count)] + if not load_weights: + for _ in e.map(self.model_comm.process_weight, self.gpu_list, ranks): + pass + if version.parse(lmdeploy.__version__) < version.parse('0.7.2'): + for _ in e.map(self.model_comm.create_engine, self.gpu_list, ranks, repeat(self.nccl_params)): + pass + else: + for _ in e.map(self.model_comm.create_engine, self.gpu_list, ranks): + pass + + def _create_weight(self, model_comm): + """Allocate weight buffer, load params if from_workspace.""" + + # TODO: support mpi + self.node_id = 0 + self.node_num = 1 + if version.parse(lmdeploy.__version__) < version.parse('0.7.2'): + self.nccl_params = model_comm.create_nccl_params(self.node_id) + synchronize() + + # create weight + def _create_weight_func(index, device_id): + rank = self.node_id * self.gpu_count + index + model_comm.create_shared_weights(device_id, rank) + + with ThreadPoolExecutor(max_workers=self.gpu_count) as executor: + futures = [] + for idx, device_id in enumerate(self.gpu_list): + futures.append(executor.submit(_create_weight_func, idx, device_id)) + for future in futures: + future.result() + + def _get_model_params(self, model_comm, tm_params): + """Get turbomind model params when loading from hf.""" + + def _get_params(idx, device_id, que): + rank = self.node_id * self.gpu_count + idx + out = model_comm.get_params(device_id, rank) + que.put(out) + + que = Queue() + with ThreadPoolExecutor(max_workers=self.gpu_count) as executor: + futures = [] + for idx, device_id in enumerate(self.gpu_list): + futures.append(executor.submit(_get_params, idx, device_id, que)) + for future in futures: + future.result() + + for _ in range(self.gpu_count): + tensor_map = que.get() + for k, v in tensor_map.items(): + if k not in tm_params: + tm_params[k] = [] + tm_params[k].append(v) + + def _load_weights(self, state_dict): + tm_params = self.tm_model.tm_params + self._get_model_params(self.model_comm, tm_params) + input_model = self.tm_model.input_model + model_path = input_model.model_path + input_model.model_path = state_dict + self.tm_model.export() + input_model.model_path = model_path + + from lmdeploy.turbomind.turbomind import TurboMindInstance + + def create_instance(self, cuda_stream_id=0): + return TurboMindInstance(self, self.config, cuda_stream_id, self.gpu_list) + + TurboMind.__origin_init__ = TurboMind.__init__ + TurboMind.__init__ = __init__ + TurboMind._create_weight = _create_weight + TurboMind._get_model_params = _get_model_params + TurboMind.create_instance = create_instance + if load_weights: + TurboMind.load_weights = _load_weights + + def __init_ins__(self, tm_model, config, cuda_stream_id=0, gpu_list=None): + if gpu_list is None: + gpu_list = [0] + self.gpu_list = gpu_list + self.__origin_init__(tm_model, config, cuda_stream_id) + + def _create_model_instance(self, device_id): + model_inst = self.tm_model.model_comm.create_model_instance(self.gpu_list[0]) + return model_inst + + TurboMindInstance.__origin_init__ = TurboMindInstance.__init__ + TurboMindInstance.__init__ = __init_ins__ + TurboMindInstance._create_model_instance = _create_model_instance + + +def patch_npu_vllm(vllm_device: str, *, colocate: bool = False): + if isinstance(vllm_device, int): + vllm_device = get_device(vllm_device) + device_type = vllm_device.split(':')[0] + if device_type == 'npu': + from swift.model.npu_patch.vllm_ascend import patch_vllm_ascend_runtime + from swift.model.npu_patch.vllm_ascend_memory import vllm_ascend_mem_get_info_context + patch_vllm_ascend_runtime(colocate=colocate) + return vllm_ascend_mem_get_info_context(vllm_device) + + return nullcontext() + + +def patch_vllm_triton_device_guard(): + import functools + try: + from vllm.v1.worker import gpu_worker as _gw + _orig_fn = _gw.init_worker_distributed_environment + except (ImportError, AttributeError): + return + + if getattr(_gw, '_swift_dist_env_patched', False): + return + + @functools.wraps(_orig_fn) + def _patched_init_worker_distributed_environment(*args, **kwargs): + if not torch.cuda.is_available(): + return _orig_fn(*args, **kwargs) + expected_device = torch.cuda.current_device() + result = _orig_fn(*args, **kwargs) + actual_device = torch.cuda.current_device() + if actual_device != expected_device: + torch.cuda.set_device(expected_device) + return result + + _gw.init_worker_distributed_environment = _patched_init_worker_distributed_environment + _gw._swift_dist_env_patched = True + + +def patch_vllm_memory_leak(): + # fix vllm 0.7.3 memory leak + # https://github.com/vllm-project/vllm/pull/14326 + import vllm + try: + vllm_version = version.parse(vllm.__version__) + needs_patch = (vllm_version == version.parse('0.7.3')) + except version.InvalidVersion: + needs_patch = False + + if not needs_patch: + return + + def patch_vllm_abort_seq_group(): + from typing import Dict, Iterable + from vllm.core.scheduler import Scheduler + from vllm.sequence import SequenceGroup, SequenceGroupBase, SequenceStatus + + def new_abort_seq_group( + self, + request_id: Union[str, Iterable[str]], + seq_id_to_seq_group: Optional[Dict[str, SequenceGroupBase]] = None, + ) -> None: + if isinstance(request_id, str): + request_id = (request_id, ) + request_ids = set(request_id) + seq_id_to_seq_group = seq_id_to_seq_group or {} + for state_queue in [self.waiting, self.running, self.swapped]: + aborted_groups: List[SequenceGroup] = [] + for seq_group in state_queue: + # When n>1, seq_group.request_id looks like + # foo_parallel_sample_0, while request_ids is just foo, and we + # should resolve it as real_request_id to match. + if seq_group.request_id in seq_id_to_seq_group: + real_request_id = seq_id_to_seq_group[seq_group.request_id].group_id + else: + real_request_id = seq_group.request_id + if real_request_id in request_ids: + # Appending aborted group into pending list. + aborted_groups.append(seq_group) + # We can't remove real_request_id in request_ids here, + # because there may be other seq groups sharing the same + # real_request_id + for aborted_group in aborted_groups: + # Remove the sequence group from the state queue. + state_queue.remove(aborted_group) + # Remove the aborted request from the Mamba cache. + self._finished_requests_ids.append(aborted_group.request_id) + for seq in aborted_group.get_seqs(): + if seq.is_finished(): + continue + seq.status = SequenceStatus.FINISHED_ABORTED + self.free_seq(seq) + if aborted_group.request_id in seq_id_to_seq_group: + del seq_id_to_seq_group[aborted_group.request_id] + + self._free_seq_group_cross_attn_blocks(aborted_group) + + origin_method = Scheduler.abort_seq_group + Scheduler._old_abort_seq_group = origin_method + Scheduler.abort_seq_group = new_abort_seq_group + + def patch_vllm_engine(): + from vllm.engine.llm_engine import LLMEngine, SchedulerOutputState + from vllm.outputs import PoolingRequestOutput, RequestOutput + from vllm.sequence import ExecuteModelRequest + + def new_abort_request(self, request_id) -> None: + for scheduler in self.scheduler: + scheduler.abort_seq_group(request_id, seq_id_to_seq_group=self.seq_id_to_seq_group) + + origin_method = LLMEngine.abort_request + LLMEngine._old_abort_request = origin_method + LLMEngine.abort_request = new_abort_request + + def new_step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]: + if self.parallel_config.pipeline_parallel_size > 1: + raise NotImplementedError('Pipeline parallelism is only supported through AsyncLLMEngine ' + 'as performance will be severely degraded otherwise.') + + # For llm_engine, there is no pipeline parallel support, so the engine + # used is always 0. + virtual_engine = 0 + + # These are cached outputs from previous iterations. None if on first + # iteration + cached_outputs = self.cached_scheduler_outputs[virtual_engine] + seq_group_metadata_list = cached_outputs.seq_group_metadata_list + scheduler_outputs = cached_outputs.scheduler_outputs + allow_async_output_proc = cached_outputs.allow_async_output_proc + + ctx = self.scheduler_contexts[virtual_engine] + + # Clear outputs for each new scheduler iteration + ctx.request_outputs.clear() + + # Skip the scheduler if there are any remaining steps in the seq groups. + # This ensures that the scheduler is only called again when the current + # batch has completed. + # The scheduler is also skipped if a single request caused the last + # engine step to fail, and the previous schedule needs to be rerun. + if not self._has_remaining_steps(seq_group_metadata_list): + # Schedule iteration + (seq_group_metadata_list, scheduler_outputs, + allow_async_output_proc) = self.scheduler[virtual_engine].schedule() + + ctx.seq_group_metadata_list = seq_group_metadata_list + ctx.scheduler_outputs = scheduler_outputs + + finished_requests_ids = self.scheduler[virtual_engine].get_and_reset_finished_requests_ids() + # When n>1, elements in self.seq_id_to_seq_group should be deleted + # here, otherwise memory leaks. + for finished_request_id in finished_requests_ids: + if finished_request_id in self.seq_id_to_seq_group: + del self.seq_id_to_seq_group[finished_request_id] + + # Maybe switch from async mode to sync mode + if not allow_async_output_proc and len(ctx.output_queue) > 0: + self._process_model_outputs(ctx=ctx) + + if (self.scheduler_config.is_multi_step and scheduler_outputs.num_lookahead_slots > 0): + # cache the scheduler outputs for the next iteration if we have + # lookahead slots + self._cache_scheduler_outputs_for_multi_step(virtual_engine, seq_group_metadata_list, + scheduler_outputs, allow_async_output_proc) + else: + finished_requests_ids = list() + + assert seq_group_metadata_list is not None + assert scheduler_outputs is not None + + if not scheduler_outputs.is_empty(): + + # Check if we have a cached last_output from the previous iteration. + # For supporting PP this is probably the best way to pass the + # sampled_token_ids, as a separate broadcast over all the PP stages + # will cause one virtual engine's microbatch to block the pipeline. + last_sampled_token_ids = \ + self._get_last_sampled_token_ids(virtual_engine) + + execute_model_req = ExecuteModelRequest( + seq_group_metadata_list=seq_group_metadata_list, + blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in, + blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out, + blocks_to_copy=scheduler_outputs.blocks_to_copy, + num_lookahead_slots=scheduler_outputs.num_lookahead_slots, + running_queue_size=scheduler_outputs.running_queue_size, + finished_requests_ids=finished_requests_ids, + # We use ExecuteModelRequest to pass the last sampled_token_ids + # to each of the non-last PP stages for in-place prepare_input. + last_sampled_token_ids=last_sampled_token_ids) + + if allow_async_output_proc: + execute_model_req.async_callback = self.async_callbacks[virtual_engine] + + outputs = self.model_executor.execute_model(execute_model_req=execute_model_req) + + # We need to do this here so that last step's sampled_token_ids can + # be passed to the next iteration for PP. + if self.scheduler_config.is_multi_step: + self._update_cached_scheduler_output(virtual_engine, outputs) + else: + # Nothing scheduled => If there is pending async postprocessor, + # then finish it here. + if len(ctx.output_queue) > 0: + self._process_model_outputs(ctx=ctx) + # No outputs in this case + outputs = [] + + # Finish the current step for all the sequence groups. + if self.scheduler_config.is_multi_step: + for seq_group in seq_group_metadata_list: + seq_group.finish_step() + + if not self._has_remaining_steps(seq_group_metadata_list): + # clear the cache if we have finished all the steps. + if self.scheduler_config.is_multi_step: + self.cached_scheduler_outputs[0] = SchedulerOutputState() + + # is_first_step_output is True only when the num_steps of all + # the sequences are 1. When the num_steps > 1, + # multi_step_model_runner does the first-step output append. + is_first_step_output: bool = False if not seq_group_metadata_list \ + else seq_group_metadata_list[0].state.num_steps == 1 + + # Add results to the output_queue + ctx.append_output( + outputs=outputs, + seq_group_metadata_list=seq_group_metadata_list, + scheduler_outputs=scheduler_outputs, + is_async=allow_async_output_proc, + is_last_step=True, + is_first_step_output=is_first_step_output) + + if outputs and allow_async_output_proc: + assert len(outputs) == 1, ('Async postprocessor expects only a single output set') + + self._advance_to_next_step(outputs[0], seq_group_metadata_list, + scheduler_outputs.scheduled_seq_groups) + + # Check if need to run the usual non-async path + if not allow_async_output_proc: + self._process_model_outputs(ctx=ctx) + + # Log stats. + self.do_log_stats(scheduler_outputs, outputs) + + # Tracing + self.do_tracing(scheduler_outputs) + else: + # Multi-step case + return ctx.request_outputs + + if not self.has_unfinished_requests(): + # Drain async postprocessor (if exists) + if len(ctx.output_queue) > 0: + self._process_model_outputs(ctx=ctx) + assert len(ctx.output_queue) == 0 + + # Stop the execute model loop in parallel workers until there are + # more requests to process. This avoids waiting indefinitely in + # torch.distributed ops which may otherwise timeout, and unblocks + # the RPC thread in the workers so that they can process any other + # queued control plane messages, such as add/remove lora adapters. + self.model_executor.stop_remote_worker_execution_loop() + + return ctx.request_outputs + + origin_method = LLMEngine.step + LLMEngine._old_step = origin_method + LLMEngine.step = new_step + + patch_vllm_abort_seq_group() + patch_vllm_engine() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/vllm_engine.py b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/vllm_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..bae239ccc9a09aec8e6176bded9b4c093615336c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/infer_engine/vllm_engine.py @@ -0,0 +1,958 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import asyncio +import inspect +import multiprocessing +import os +import time +import torch +from contextlib import contextmanager, nullcontext +from copy import copy, deepcopy +from packaging import version +from PIL import Image +from tqdm import tqdm +from transformers import AutoConfig, GenerationConfig +from transformers.utils import is_torch_npu_available +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union + +from swift.metrics import Metric +from swift.model import get_processor +from swift.template import Template +from swift.utils import (HfConfigFactory, disable_deepspeed_zero3, get_device, get_dist_setting, get_logger, is_dist, + safe_snapshot_download) +from .infer_engine import InferEngine +from .patch import patch_auto_tokenizer +from .protocol import (ChatCompletionResponse, ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, ChatMessage, DeltaMessage, EmbeddingResponse, + EmbeddingResponseData, InferRequest, RequestConfig, random_uuid) +from .utils import AdapterRequest, InferStreamer, patch_npu_vllm, patch_vllm_memory_leak, patch_vllm_triton_device_guard + +logger = get_logger() +try: + # After setting the environment variables, import vllm. This way of writing allows lint to pass. + os.environ['VLLM_WORKER_MULTIPROC_METHOD'] = 'spawn' + os.environ['VLLM_ENGINE_ITERATION_TIMEOUT_S'] = '86400' + import vllm + from vllm import AsyncEngineArgs, AsyncLLMEngine, EngineArgs, LLMEngine, SamplingParams + from vllm.pooling_params import PoolingParams + try: + # vLLM v0.12+ uses StructuredOutputsParams + from vllm.sampling_params import StructuredOutputsParams + except ImportError: + # Fallback for older vLLM versions + from vllm.sampling_params import GuidedDecodingParams as StructuredOutputsParams +except Exception: + raise + +try: + from vllm.reasoning import ReasoningParserManager +except ImportError: + ReasoningParserManager = None + +dtype_mapping = {torch.float16: 'float16', torch.bfloat16: 'bfloat16', torch.float32: 'float32'} + + +def _patch_vllm_dp_coordinator_timeout(): + # https://github.com/vllm-project/vllm/pull/37452 introduced a 30-second default timeout, + # which is prone to timing out in spawn scenarios. Patch it to 180 seconds here. + try: + from vllm.v1.engine import coordinator as coordinator_module + except ImportError: + return + + coordinator_cls = coordinator_module.DPCoordinator + if not hasattr(coordinator_cls, '_wait_for_zmq_addrs'): + return + + if getattr(coordinator_cls, '_swift_timeout_patched', False): + return + + def _wait_for_zmq_addrs(self, zmq_addr_pipe): + t0 = time.monotonic() + try: + ready = multiprocessing.connection.wait([zmq_addr_pipe, self.proc.sentinel], timeout=180) + elapsed = time.monotonic() - t0 + if not ready: + raise RuntimeError(f'DP Coordinator process failed to report ZMQ addresses ' + f'within 180s (elapsed={elapsed:.1f}s).') + try: + return zmq_addr_pipe.recv() + except EOFError: + raise RuntimeError('DP Coordinator process failed during startup.') from None + finally: + zmq_addr_pipe.close() + + coordinator_cls._wait_for_zmq_addrs = _wait_for_zmq_addrs + coordinator_cls._swift_timeout_patched = True + + +_patch_vllm_dp_coordinator_timeout() + + +@contextmanager +def _patch_rope_validation_ignore_keys(): + """Accept list-style RoPE validation ignore keys from older vLLM configs. + + vLLM 0.18.x Qwen3.5 configs may pass ``ignore_keys_at_rope_validation`` + as a list, while Transformers 5.x treats it as a set and performs a set + union during RoPE validation. vLLM release tags from 0.19.0 onward changed + the Qwen3.5 configs to set literals, but 0.18-based vLLM/vLLM-Ascend stacks + still need this compatibility layer. See vLLM PR: + https://github.com/vllm-project/vllm/pull/37338 + """ + from transformers import PretrainedConfig + + origin_convert = getattr(PretrainedConfig, 'convert_rope_params_to_dict', None) + if origin_convert is None: + yield + return + + def convert_rope_params_to_dict(self, ignore_keys_at_rope_validation=None, **kwargs): + if isinstance(ignore_keys_at_rope_validation, list): + ignore_keys_at_rope_validation = set(ignore_keys_at_rope_validation) + return origin_convert(self, ignore_keys_at_rope_validation=ignore_keys_at_rope_validation, **kwargs) + + PretrainedConfig.convert_rope_params_to_dict = convert_rope_params_to_dict + try: + yield + finally: + PretrainedConfig.convert_rope_params_to_dict = origin_convert + + +class VllmEngine(InferEngine): + + def __init__( + self, + model_id_or_path: str, + *, + template: Optional[Template] = None, + torch_dtype: Optional[torch.dtype] = None, + adapters: Optional[List[str]] = None, + use_async_engine: bool = False, + model_type: Optional[str] = None, + template_type: Optional[str] = None, + use_hf: Optional[bool] = None, + hub_token: Optional[str] = None, + revision: Optional[str] = None, + # engine_kwargs + gpu_memory_utilization: float = 0.9, + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + enable_expert_parallel: bool = False, + max_model_len: Optional[int] = None, + max_num_seqs: int = 256, + disable_custom_all_reduce: bool = True, + enforce_eager: bool = False, + limit_mm_per_prompt: Optional[Dict[str, Any]] = None, + seed: Optional[int] = None, + task_type: Optional[str] = None, # embedding + disable_cascade_attn: bool = False, + load_format: str = 'auto', + mm_processor_cache_gb: Optional[float] = None, + logprobs_mode: Optional[str] = None, + speculative_config: Optional[Union[str, dict]] = None, + # lora + enable_lora: bool = False, + max_loras: int = 1, + max_lora_rank: int = 16, + enable_prefix_caching: Optional[bool] = None, + enable_sleep_mode: bool = False, + distributed_executor_backend: Optional[str] = None, + quantization: Optional[str] = None, + # reasoning parser + reasoning_parser: Optional[str] = None, + engine_kwargs: Optional[Dict[str, Any]] = None, + num_labels: Optional[int] = None, + reranker_use_activation: bool = True, + ) -> None: + self.model_id_or_path = model_id_or_path + self.torch_dtype = torch_dtype + if isinstance(adapters, str): + adapters = [adapters] + self.default_adapter_request = None + if isinstance(adapters, list) and adapters: + assert len(adapters) == 1, 'Only one adapter is supported for now.' + enable_lora = True + self.default_adapter_request = AdapterRequest('default', adapters[0]) + self.adapters = adapters or [] + self.use_async_engine = use_async_engine + self.model_type = model_type + self.use_hf = use_hf + self.hub_token = hub_token + self.revision = revision + + self.gpu_memory_utilization = gpu_memory_utilization + self.tensor_parallel_size = tensor_parallel_size + self.pipeline_parallel_size = pipeline_parallel_size + self.enable_expert_parallel = enable_expert_parallel + self.max_num_seqs = max_num_seqs + self.disable_custom_all_reduce = disable_custom_all_reduce + self.enforce_eager = enforce_eager + self.limit_mm_per_prompt = limit_mm_per_prompt + self.seed = seed + self.task_type = task_type + self.disable_cascade_attn = disable_cascade_attn + self.load_format = load_format + self.mm_processor_cache_gb = mm_processor_cache_gb + self.logprobs_mode = logprobs_mode + self.speculative_config = speculative_config + + self.enable_lora = enable_lora + self.max_loras = max_loras + self.max_lora_rank = max_lora_rank + self.enable_prefix_caching = enable_prefix_caching + self.enable_sleep_mode = enable_sleep_mode + self.distributed_executor_backend = distributed_executor_backend + self.quantization = quantization + self.num_labels = num_labels + self.reranker_use_activation = reranker_use_activation + + patch_vllm_memory_leak() + patch_vllm_triton_device_guard() + self._adapters_pool = {} + if template is None: + processor = self._get_processor() + template = self._get_template(processor, template_type=template_type) + else: + safe_snapshot_download( + model_id_or_path, + revision=revision, + download_model=True, + use_hf=use_hf, + ignore_patterns=getattr(template.model_meta, 'ignore_patterns', None), + hub_token=hub_token) + super().__init__(template) + if max_model_len is not None: + self.max_model_len = max_model_len + logger.info(f'Setting max_model_len: {max_model_len}') + self._prepare_engine_kwargs(max_model_len, engine_kwargs) + context = nullcontext() + if is_torch_npu_available() and (tensor_parallel_size == 1 or pipeline_parallel_size == 1): + colocate = ( + getattr(self, '_swift_vllm_colocate_runtime', False) + or self.distributed_executor_backend == 'external_launcher') + context = patch_npu_vllm(get_device(), colocate=colocate) + with context: + self._prepare_engine() + self._load_generation_config() + self._fix_vllm_bug() + self.patch_remove_log() + self._request_count = 0 + self._prepare_reasoning_parser(reasoning_parser) + + def _get_processor(self): + return get_processor( + model_id_or_path=self.model_id_or_path, + torch_dtype=self.torch_dtype, + download_model=True, + model_type=self.model_type, + use_hf=self.use_hf, + hub_token=self.hub_token, + revision=self.revision, + num_labels=self.num_labels, + task_type=self.task_type) + + def _prepare_engine(self) -> None: + with patch_auto_tokenizer(self.tokenizer), \ + _patch_rope_validation_ignore_keys(), disable_deepspeed_zero3(): + llm_engine_cls = AsyncLLMEngine if self.use_async_engine else LLMEngine + engine = llm_engine_cls.from_engine_args(self.engine_args) + self.engine = engine + + def _get_hf_config_overrides(self) -> Dict[str, Any]: + overrides = {} + config = self.config + try: + disk_config = AutoConfig.from_pretrained( + self.model_dir, trust_remote_code=self.model_meta.loader.default_trust_remote_code) + except Exception: + disk_config = None + + def _changed(key): + value = HfConfigFactory.get_config_attr(config, key) + if value is None or disk_config is None: + return None + return value if value != HfConfigFactory.get_config_attr(disk_config, key) else None + + rope_key = 'rope_parameters' if hasattr(config, 'rope_parameters') else 'rope_scaling' + rope_scaling = _changed(rope_key) + if rope_scaling: + rope_scaling = dict(rope_scaling) + if 'rope_type' not in rope_scaling and 'type' in rope_scaling: + rope_scaling['rope_type'] = rope_scaling['type'] + overrides[rope_key] = rope_scaling + vocab_size = _changed('vocab_size') + if vocab_size is not None: + overrides['vocab_size'] = vocab_size + if self.task_type in {'seq_cls', 'reranker'}: + for key in ['num_labels', 'problem_type']: + value = getattr(config, key, None) + if value is not None: + overrides[key] = value + return overrides + + def _prepare_engine_kwargs(self, max_model_len, engine_kwargs) -> None: + if engine_kwargs is None: + engine_kwargs = {} + encode_task_mapping = { + 'embedding': 'embed', + 'seq_cls': 'classify', + 'reranker': 'score', + 'generative_reranker': 'score', + } + vllm_task = encode_task_mapping.get(self.task_type) + disable_log_stats = engine_kwargs.pop('disable_log_stats', True) + if self.use_async_engine: + engine_cls = AsyncEngineArgs + else: + engine_cls = EngineArgs + parameters = inspect.signature(engine_cls).parameters + if vllm_task: + if 'runner' in parameters: + engine_kwargs['runner'] = 'pooling' + elif 'task' in parameters: + engine_kwargs['task'] = vllm_task + else: + raise ValueError( + f'task_type={self.task_type} requires a vLLM version that supports `runner` or `task`. ' + 'Please upgrade vLLM.') + if self.use_async_engine and 'disable_log_requests' in parameters: + engine_kwargs['disable_log_requests'] = True + if 'enable_lora' in parameters and self.enable_lora: + engine_kwargs['enable_lora'] = self.enable_lora + engine_kwargs['max_loras'] = self.max_loras + engine_kwargs['max_lora_rank'] = self.max_lora_rank + else: + assert not self.enable_lora, ( + 'The current version of vLLM does not support `enable_lora`. Please upgrade vLLM.') + + if 'limit_mm_per_prompt' in parameters and self.limit_mm_per_prompt: + engine_kwargs['limit_mm_per_prompt'] = self.limit_mm_per_prompt + else: + assert not self.limit_mm_per_prompt, ( + 'The current version of vLLM does not support `limit_mm_per_prompt`. Please upgrade vLLM.') + for key in [ + 'enable_expert_parallel', 'enable_sleep_mode', 'disable_cascade_attn', 'load_format', + 'mm_processor_cache_gb', 'speculative_config', 'logprobs_mode', 'quantization' + ]: + if key in parameters: + value = getattr(self, key, None) + if value is not None: + engine_kwargs[key] = value + else: + logger.warning(f'The current version of vLLM does not support `{key}`. Ignored.') + if self.seed is not None: + engine_kwargs['seed'] = self.seed + + model_info = self.model_info + hf_overrides = engine_kwargs.pop('hf_overrides', None) or {} + arch_mapping = {'deepseek_vl2': ['DeepseekVLV2ForCausalLM'], 'chatglm4v': ['GLM4VForCausalLM']} + if self.model_meta.model_type in arch_mapping: + hf_overrides['architectures'] = arch_mapping[self.model_meta.model_type] + hf_overrides.update(self._get_hf_config_overrides()) + if hf_overrides: + engine_kwargs['hf_overrides'] = hf_overrides + self.template.set_mode('vllm') + engine_kwargs.update(self.template.prepare_engine_kwargs()) + if self.enable_prefix_caching is not None: + engine_kwargs['enable_prefix_caching'] = self.enable_prefix_caching + engine_args = engine_cls( + model=self.model_dir, + dtype=dtype_mapping[model_info.torch_dtype], + gpu_memory_utilization=self.gpu_memory_utilization, + tensor_parallel_size=self.tensor_parallel_size, + pipeline_parallel_size=self.pipeline_parallel_size, + max_model_len=max_model_len, + max_num_seqs=self.max_num_seqs, + disable_log_stats=disable_log_stats, + disable_custom_all_reduce=self.disable_custom_all_reduce, + enforce_eager=self.enforce_eager, + trust_remote_code=True, + distributed_executor_backend=self.distributed_executor_backend, + **engine_kwargs, + ) + self.engine_args = engine_args + + def _prepare_reasoning_parser(self, reasoning_parser: Optional[str]) -> None: + self.reasoning_parser = None + if not reasoning_parser: + return + + # Validate reasoning_parser if provided + if ReasoningParserManager is None: + raise ImportError('the version of vLLM is too old, please upgrade vLLM') + + valid_reasoning_parsers = list(ReasoningParserManager.reasoning_parsers.keys()) + if reasoning_parser not in valid_reasoning_parsers: + raise ValueError(f'Invalid reasoning_parser: {reasoning_parser}. ' + f'Available parsers: {valid_reasoning_parsers}') + logger.info(f'Using reasoning_parser: {reasoning_parser}') + + reasoning_parser_cls = ReasoningParserManager.get_reasoning_parser(reasoning_parser) + self.reasoning_parser = reasoning_parser_cls(self.tokenizer) + + def _fix_vllm_bug(self) -> None: + # fix vllm==0.4 bug (very slow) + tokenizer = self.tokenizer + if self._version_ge( + '0.4') and not self._version_ge('0.6') and not tokenizer.__class__.__name__.startswith('Cached'): + _tokenizer_len = len(tokenizer) + __old_len__ = tokenizer.__class__.__len__ + + def __len__(self) -> int: + if self is tokenizer: + return _tokenizer_len + else: + return __old_len__(self) + + tokenizer.__class__.__len__ = __len__ + + def _load_generation_config(self) -> None: + generation_config_path = os.path.join(self.model_dir, 'generation_config.json') + if os.path.isfile(generation_config_path): + generation_config = GenerationConfig.from_pretrained(self.model_dir) + kwargs = generation_config.to_dict() + max_new_tokens = kwargs.get('max_new_tokens') + if max_new_tokens is not None: + kwargs['max_tokens'] = max_new_tokens + top_k = kwargs.get('top_k') + if top_k == 0: + kwargs['top_k'] = -1 + parameters = inspect.signature(SamplingParams).parameters + for k, v in kwargs.copy().items(): + if k not in parameters or v is None: + kwargs.pop(k) + self.generation_config = SamplingParams(**kwargs) + else: + self.generation_config = SamplingParams() + + def _add_stop_words(self, generation_config: SamplingParams, request_config: RequestConfig) -> None: + template_meta = self.template.template_meta + stop_words = (request_config.stop or []) + (self.generation_config.stop or []) + template_meta.stop_words + generation_config.stop = self._get_stop_words(stop_words) + # stop parameter is not effective in v1 engine (test version: vllm 0.8.5.post) + generation_config.stop_token_ids = self._get_stop_token_ids(stop_words) + + @staticmethod + def _version_ge(base_version: str): + vllm_version = vllm.__version__ + if vllm_version is None or 'dev' in vllm_version: + return True + return version.parse(vllm_version) >= version.parse(base_version) + + def _add_adapter(self, adapter_request: Optional[AdapterRequest] = None): + assert self.enable_lora, f'adapter_request: {adapter_request}, self.enable_lora: {self.enable_lora}' + from vllm.lora.request import LoRARequest + adapter_name = adapter_request.name + adapter_path = adapter_request.path + if adapter_name in self._adapters_pool: + lora_request = self._adapters_pool[adapter_name] + else: + lora_request = LoRARequest( + lora_name=adapter_name, lora_path=adapter_path, lora_int_id=len(self._adapters_pool) + 1) + self._adapters_pool[adapter_name] = lora_request + return lora_request + + def _add_request(self, + inputs: Dict[str, Any], + generation_config: SamplingParams, + request_id: str, + adapter_request: Optional[AdapterRequest] = None): + kwargs = {} + adapter_request = adapter_request or self.default_adapter_request + if adapter_request: + kwargs['lora_request'] = self._add_adapter(adapter_request) + + input_ids = inputs['input_ids'] + if self._version_ge('0.4.3'): + llm_inputs = {'prompt_token_ids': input_ids} + mm_data = {} + for key in ['images', 'audios', 'videos']: + media_data = inputs.get(key) or [] + if media_data: + if self._version_ge('0.6'): + + mm_data[key.rstrip('s')] = media_data[0] if ( + len(media_data) == 1 and + # compat qwen3_vl + not isinstance(media_data[0], tuple)) else media_data + else: + assert len(media_data) == 1, ( + f'The current version of vllm only supports single {key}. Please upgrade to vllm >= 0.6.0') + mm_data[key.rstrip('s')] = media_data[0] + if mm_data: + llm_inputs['multi_modal_data'] = mm_data + mm_processor_kwargs = inputs.get('mm_processor_kwargs') + if mm_processor_kwargs: + llm_inputs['mm_processor_kwargs'] = mm_processor_kwargs + + has_task_arg = 'task' in inspect.signature(PoolingParams).parameters + has_activation_arg = 'activation' in inspect.signature(PoolingParams).parameters + task_mapping = { + 'embedding': 'embed', + 'seq_cls': 'classify', + 'reranker': 'score', + 'generative_reranker': 'score', + } + if self.task_type in task_mapping: + pooling_kwargs = {} + if has_task_arg: + pooling_kwargs['task'] = task_mapping[self.task_type] + if self.task_type in ('reranker', 'generative_reranker') and \ + has_activation_arg and self.reranker_use_activation: + pooling_kwargs['activation'] = True + pooling_params = PoolingParams(**pooling_kwargs) + return self.engine.encode(llm_inputs, pooling_params, request_id) + elif self.use_async_engine: + return self.engine.generate(llm_inputs, generation_config, request_id, **kwargs) + else: + return self.engine.add_request(request_id, llm_inputs, generation_config, **kwargs) + else: + if self.use_async_engine: + return self.engine.generate(None, generation_config, request_id, input_ids, **kwargs) + else: + return self.engine.add_request(request_id, None, generation_config, input_ids, **kwargs) + + def _get_logprobs(self, + logprobs_list: Optional[List[Dict[int, float]]], + token_ids: List[int], + top_logprobs: Optional[int] = None) -> Optional[Dict[str, Any]]: + if logprobs_list is None or len(token_ids) == 0: + return None + if len(token_ids) > 0: + logprobs_list = logprobs_list[-len(token_ids):] + for logprobs in logprobs_list: + for token_id, logprob in logprobs.items(): + logprobs[token_id] = logprob.logprob + return super()._get_logprobs(logprobs_list, token_ids, top_logprobs) + + def _prepare_generation_config(self, request_config: RequestConfig) -> SamplingParams: + kwargs = {'max_tokens': request_config.max_tokens} + for key in ['temperature', 'top_k', 'top_p', 'min_p', 'repetition_penalty']: + new_value = getattr(request_config, key) + if new_value is None: + kwargs[key] = getattr(self.generation_config, key) + else: + kwargs[key] = new_value + + # Convert Swift's Chat Completions API style (logprobs: bool, top_logprobs: int) + # to vLLM's SamplingParams style (logprobs: int) + # vLLM semantics: + # - logprobs=None: no logprobs returned + # - logprobs=0: only sampled token's logprob + # - logprobs=N: top-N tokens + sampled token (up to N+1 total) + if request_config.logprobs: + # If logprobs=True, return log probabilities + if request_config.top_logprobs is not None and request_config.top_logprobs > 0: + # Return top_logprobs most likely tokens at each position + # (plus sampled token if not in top-N) + kwargs['logprobs'] = request_config.top_logprobs + else: + # Return only the sampled token's logprob + kwargs['logprobs'] = 0 + + if request_config.prompt_logprobs is not None: + kwargs['prompt_logprobs'] = request_config.prompt_logprobs + + # TODO: beam search + for key in ['n', 'best_of', 'frequency_penalty', 'presence_penalty', 'seed']: + if hasattr(SamplingParams, key): + kwargs[key] = getattr(request_config, key) + + # Handle structured outputs (guided decoding) + # vLLM v0.12+ uses 'structured_outputs' parameter, older versions use 'guided_decoding' + if request_config.structured_outputs_regex: + structured_outputs_param = StructuredOutputsParams(regex=request_config.structured_outputs_regex) + if hasattr(SamplingParams, 'structured_outputs'): + kwargs['structured_outputs'] = structured_outputs_param + else: + # Fallback for older vLLM versions + kwargs['guided_decoding'] = structured_outputs_param + + res = SamplingParams(**kwargs) + + if hasattr(res, 'output_kind') and res.n > 1: + # fix n > 1 in V1 Engine + from vllm.sampling_params import RequestOutputKind + res.output_kind = RequestOutputKind.FINAL_ONLY + return res + + @property + def inner_model(self): + return self.engine.model_executor.driver_worker.worker.model_runner.model + + @property + def inner_model_executor(self): + return self.engine.model_executor + + async def _infer_stream_async( + self, + inputs: Dict[str, Any], + generation_config: SamplingParams, + adapter_request: Optional[AdapterRequest], + request_config: RequestConfig, + ) -> AsyncIterator[ChatCompletionStreamResponse]: + request_id = random_uuid() + result_generator = self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request) + infer_streamers = [ + InferStreamer(self.template, template_inputs=inputs['template_inputs']) for _ in range(generation_config.n) + ] + token_idxs = [0 for _ in range(generation_config.n)] + async for result in result_generator: + res = self._create_chat_completion_stream_response(result, request_config, request_id, infer_streamers, + token_idxs) + if res is None: + continue + yield res + + def _create_chat_completion_stream_response(self, result, request_config, request_id, infer_streamers, + token_idxs) -> Optional[ChatCompletionStreamResponse]: + is_diff = False + is_finished = False + for output in result.outputs: + output.token_ids = list(output.token_ids) + output.delta_text = infer_streamers[output.index].get_printable_text(output.token_ids, output.finished()) + output.is_finished = output.finish_reason is not None + is_diff |= bool(output.delta_text) + is_finished |= output.is_finished + if not is_diff and not is_finished: + return + + num_generated_tokens = sum(len(output.token_ids) for output in result.outputs) + usage_info = self._get_usage_info(len(result.prompt_token_ids), num_generated_tokens) + choices = [] + previous_texts = [''] * len(result.outputs) + for output in result.outputs: + i = output.index + logprobs = self._get_logprobs(output.logprobs, output.token_ids[token_idxs[i]:], + request_config.top_logprobs) + + # Handle reasoning content in streaming + delta_content = output.delta_text + delta_reasoning_content = None + + if self.reasoning_parser and output.delta_text: + try: + # Get token IDs for the delta (new tokens in this step) + delta_token_ids = output.token_ids[token_idxs[i]:] + previous_token_ids = output.token_ids[:token_idxs[i]] + + # Get current accumulated text for this output + previous_text = previous_texts[i] + current_text = previous_text + output.delta_text + previous_texts[i] = current_text + # Extract reasoning content from the delta + delta_message = self.reasoning_parser.extract_reasoning_content_streaming( + previous_text, current_text, output.delta_text, previous_token_ids, output.token_ids, + delta_token_ids) + + if delta_message: + delta_reasoning_content = delta_message.reasoning_content + if delta_message.content: + delta_content = delta_message.content + else: + delta_content = None + + except Exception as e: + logger.warning(f'Failed to extract reasoning content in streaming: {e}') + # Fallback to original delta_text + delta_content = output.delta_text + token_idxs[i] = len(output.token_ids) + + toolcall = None + if output.is_finished: + toolcall = self._get_toolcall( + self.template.decode_generate_ids(output.token_ids, **infer_streamers[i].decode_kwargs)) + + choice = ChatCompletionResponseStreamChoice( + index=i, + delta=DeltaMessage( + role='assistant', + content=delta_content, + reasoning_content=delta_reasoning_content, + tool_calls=toolcall), + finish_reason=output.finish_reason, + logprobs=logprobs) + choices.append(choice) + return ChatCompletionStreamResponse(model=self.model_name, choices=choices, usage=usage_info, id=request_id) + + @staticmethod + def _format_prompt_logprobs(prompt_logprobs): + if prompt_logprobs is None: + return None + result = [] + for pos_lps in prompt_logprobs: + if pos_lps is None: + result.append(None) + else: + pos_dict = {} + for token_id, lp_obj in pos_lps.items(): + pos_dict[str(token_id)] = { + 'logprob': lp_obj.logprob, + 'rank': getattr(lp_obj, 'rank', None), + 'decoded_token': getattr(lp_obj, 'decoded_token', ''), + } + result.append(pos_dict) + return result + + def _create_embedding_response(self, result, generation_config, request_id) -> EmbeddingResponse: + assert result is not None + embedding = result.outputs.data.cpu().numpy().tolist() + usage_info = self._get_usage_info(len(result.prompt_token_ids), 0) + return EmbeddingResponse( + model=self.model_name, data=[EmbeddingResponseData(embedding=embedding)], usage=usage_info, id=request_id) + + def _create_chat_completion_response( + self, + result, + inputs, + request_config, + request_id, + ) -> ChatCompletionResponse: + assert result is not None + num_generated_tokens = sum(len(output.token_ids) for output in result.outputs) + usage_info = self._get_usage_info(len(result.prompt_token_ids), num_generated_tokens) + choices = [] + for output in result.outputs: + output.token_ids = list(output.token_ids) + response = self.template.decode_generate_ids(output.token_ids, template_inputs=inputs['template_inputs']) + + # Extract reasoning content if reasoning_parser is enabled + reasoning_content = None + content = response + if self.reasoning_parser: + try: + reasoning_content, content = self.reasoning_parser.extract_reasoning_content( + response, + request=None # We don't have the original request here + ) + except Exception as e: + logger.warning(f'Failed to extract reasoning content: {e}') + # Fallback to original response + content = response + + logprobs = self._get_logprobs(output.logprobs, output.token_ids, request_config.top_logprobs) + toolcall = self._get_toolcall(content) # Use content instead of response for tool calls + token_ids = output.token_ids if request_config.return_details else None + choice = ChatCompletionResponseChoice( + index=output.index, + message=ChatMessage( + role='assistant', content=content, reasoning_content=reasoning_content, tool_calls=toolcall), + finish_reason=output.finish_reason, + logprobs=logprobs, + token_ids=token_ids) + choices.append(choice) + prompt_token_ids = None + images_size = None + if request_config.return_details: + prompt_token_ids = result.prompt_token_ids + images = inputs['template_inputs'].images + if all(isinstance(image, Image.Image) for image in images): + images_size = [image.size for image in images] + formatted_prompt_logprobs = None + if request_config.prompt_logprobs is not None: + formatted_prompt_logprobs = self._format_prompt_logprobs(result.prompt_logprobs) + return ChatCompletionResponse( + model=self.model_name, + choices=choices, + usage=usage_info, + id=request_id, + prompt_token_ids=prompt_token_ids, + prompt_logprobs=formatted_prompt_logprobs, + images_size=images_size) + + def _create_seq_cls_response( + self, + result, + request_config, + request_id, + ) -> ChatCompletionResponse: + assert result is not None + choices = [] + preds = result.outputs.data + if preds.dim() == 1: + preds = preds.unsqueeze(0) + if self.task_type == 'seq_cls': + top_logprobs = request_config.top_logprobs or 20 + preds, logprobs = self.template.decode_seq_cls(preds, top_logprobs) + else: + logprobs = [None] * len(preds) + num_prompt_token_ids = 0 + num_generated_tokens = 0 + for i, pred in enumerate(preds): + num_prompt_token_ids += len(result.prompt_token_ids) + num_generated_tokens += 1 + if isinstance(pred, torch.Tensor): + pred = pred.tolist() + choices.append( + ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role='assistant', content=pred, tool_calls=None), + finish_reason='stop', + logprobs=logprobs[i])) + usage_info = self._get_usage_info(num_prompt_token_ids, num_generated_tokens) + return ChatCompletionResponse( + model=self.model_name, + choices=choices, + usage=usage_info, + id=request_id, + prompt_token_ids=result.prompt_token_ids) + + async def _infer_full_async( + self, + inputs: Dict[str, Any], + generation_config: SamplingParams, + adapter_request: Optional[AdapterRequest], + request_config: RequestConfig, + request_id: Optional[str] = None, + ) -> Union[ChatCompletionResponse, EmbeddingResponse]: + if request_id is None: + request_id = random_uuid() + result_generator = self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request) + result = None + async for result in result_generator: + pass + if self.task_type == 'embedding': + return self._create_embedding_response(result, generation_config, request_id) + elif self.task_type in ('seq_cls', 'reranker', 'generative_reranker'): + return self._create_seq_cls_response(result, request_config, request_id) + else: + return self._create_chat_completion_response(result, inputs, request_config, request_id) + + def _batch_infer_stream(self, *args, **kwargs): + if hasattr(self.engine, 'engine'): + self.engine.engine.model_executor.parallel_worker_tasks = None + return super()._batch_infer_stream(*args, **kwargs) + + def infer( + self, + infer_requests: List[InferRequest], + request_config: Optional[RequestConfig] = None, + metrics: Optional[List[Metric]] = None, + *, + use_tqdm: Optional[bool] = None, + adapter_request: Optional[AdapterRequest] = None, + ) -> List[Union[ChatCompletionResponse, Iterator[ChatCompletionStreamResponse]]]: + if self.use_async_engine: + return super().infer( + infer_requests, + request_config, + metrics, + use_tqdm=use_tqdm, + adapter_request=adapter_request, + ) + else: + request_config = deepcopy(request_config or RequestConfig()) + if request_config.stream and len(infer_requests) > 1: + raise ValueError('If you want to use stream batch inference, you need to set use_async_engine to True.') + if use_tqdm is None: + use_tqdm = len(infer_requests) > 1 + rank = get_dist_setting()[0] + if is_dist() and rank % self.engine_args.tensor_parallel_size != 0: + use_tqdm = False + self.template.set_mode('vllm') + batched_inputs, error_list = self._batch_encode(infer_requests, strict=getattr(self, 'strict', True)) + request_id_list = [] + for i, inputs in enumerate(batched_inputs): + request_id = str(self._request_count) + request_id_list.append(request_id) + self._request_count += 1 + _request_config = deepcopy(request_config) + self.set_default_max_tokens(_request_config, inputs) + generation_config = self._prepare_generation_config(_request_config) + if generation_config.seed is not None: + generation_config.seed += i + self._add_stop_words(generation_config, _request_config) + self._add_request(inputs, generation_config, request_id, adapter_request=adapter_request) + prog_bar = tqdm(total=len(batched_inputs), dynamic_ncols=True, disable=not use_tqdm) + outputs = {} + if request_config.stream: + + def _gen_wrapper(): + infer_streamers = [ + InferStreamer(self.template, template_inputs=inputs['template_inputs']) + for _ in range(generation_config.n) + ] + token_idxs = [0 for _ in range(generation_config.n)] + while self.engine.has_unfinished_requests(): + result = self.engine.step() + if not result: + continue + result = result[0] + res = self._create_chat_completion_stream_response(result, request_config, request_id, + infer_streamers, token_idxs) + if res is None: + continue + yield res + if result.finished: + break + + self._update_metrics(res, metrics) + + return [_gen_wrapper()] + else: + while self.engine.has_unfinished_requests(): + step_outputs = self.engine.step() + for output in step_outputs: + if output.finished: + outputs[output.request_id] = output + prog_bar.update() + prog_bar.close() + outputs = [outputs[request_id] for request_id in request_id_list] + res = [ + self._create_chat_completion_response(result, inputs, request_config, request_id) + for request_id, inputs, result in zip(request_id_list, batched_inputs, outputs) + ] + self._update_metrics(res, metrics) + return self._add_error_list(res, error_list) + + async def infer_async( + self, + infer_request: InferRequest, + request_config: Optional[RequestConfig] = None, + *, + adapter_request: Optional[AdapterRequest] = None, + pre_infer_hook=None, + ) -> Union[ChatCompletionResponse, AsyncIterator[ChatCompletionStreamResponse]]: + if not self.use_async_engine: + raise ValueError('If you want to use `infer_async`, you need to pass `use_async_engine` as True.') + request_config = deepcopy(request_config or RequestConfig()) + self.template.set_mode('vllm') + loop = asyncio.get_running_loop() + with torch.inference_mode(): + inputs = await loop.run_in_executor(None, self.template.encode, infer_request, True) + self.set_default_max_tokens(request_config, inputs) + generation_config = self._prepare_generation_config(request_config) + self._add_stop_words(generation_config, request_config) + kwargs = { + 'inputs': inputs, + 'generation_config': generation_config, + 'adapter_request': adapter_request, + 'request_config': request_config, + } + if hasattr(infer_request, 'uuid') and infer_request.uuid: + # RolloutInferRequest + kwargs.update({'request_id': infer_request.uuid}) + if pre_infer_hook: + kwargs = pre_infer_hook(kwargs) + if request_config.stream: + return self._infer_stream_async(**kwargs) + else: + return await self._infer_full_async(**kwargs) + + @staticmethod + def patch_remove_log(): + from vllm.engine import async_llm_engine + if not hasattr(async_llm_engine, '_log_task_completion'): + return + + async_llm_engine._origin_log_task_completion = async_llm_engine._log_task_completion + + def new_log_task_completion(task, error_callback) -> None: + try: + return_value = task.result() + raise AssertionError(f'The engine background task should never finish without an ' + f'exception. {return_value}') + except asyncio.exceptions.CancelledError: + pass + + async_llm_engine._log_task_completion = new_log_task_completion diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..82dbf5bdfd544231e652ef6dba3764fbefc0e289 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import BaseLoss +from .mapping import loss_map diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss/base.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss/base.py new file mode 100644 index 0000000000000000000000000000000000000000..957077ba8c83e93f97ed0d1c9fb7453990a95cf7 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss/base.py @@ -0,0 +1,53 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from swift.trainers import Trainer, TrainingArguments + + +class BaseLoss(ABC): + """Abstract base class for custom loss functions. + + This class provides a common interface for implementing custom loss functions + that can be integrated with the ms-swift training framework. All custom loss + implementations should inherit from this class and implement the __call__ method. + + Attributes: + args (TrainingArguments): Training configuration and hyperparameters. + trainer (Trainer): Reference to the trainer instance for accessing model + and training state. + """ + + def __init__(self, args: 'TrainingArguments', trainer: 'Trainer'): + """Initialize the loss function with training arguments and trainer. + + Args: + args (TrainingArguments): Training configuration and hyperparameters. + trainer (Trainer): Reference to the trainer instance. + """ + self.args = args + self.trainer = trainer + mro_class_names = [cls.__name__ for cls in trainer.__class__.__mro__] + self.is_megatron = 'BaseMegatronTrainer' in mro_class_names + + @abstractmethod + def __call__(self, outputs, labels, *, num_items_in_batch=None, loss_scale=None, **kwargs) -> torch.Tensor: + """Calculate the loss value. + + This method must be implemented by all subclasses to define the specific + loss calculation logic. + + Args: + outputs: Model outputs. + labels: Ground truth labels or targets. + num_items_in_batch (int, optional): Number of items (tokens) in the current batch, + Defaults to None. + loss_scale (float, optional): Scaling factor to apply to the loss value. + Defaults to None. + + Returns: + torch.Tensor: A scalar tensor representing the computed loss value. + """ + pass diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss/causal_lm.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss/causal_lm.py new file mode 100644 index 0000000000000000000000000000000000000000..62f802c097d1c1e9333799254bf3d9fee8e1bb1d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss/causal_lm.py @@ -0,0 +1,12 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import BaseLoss + + +class CustomCrossEntropyLoss(BaseLoss): + + def __call__(self, outputs, labels, *, num_items_in_batch=None, loss_scale=None, **kwargs): + from swift.trainers import per_token_loss_func + token_loss = per_token_loss_func(outputs, labels) + if num_items_in_batch is None: + num_items_in_batch = (labels[:, 1:] != -100).sum() + return token_loss.sum() / num_items_in_batch diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss/embedding.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..97d28f621bcdc036189b3552494ad69347c506e6 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss/embedding.py @@ -0,0 +1,321 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import numpy as np +import os +import torch +import torch.distributed as dist +import torch.nn.functional as F +from accelerate.utils import gather_object +from enum import Enum +from torch import nn +from torch.nn import MSELoss +from transformers.utils import strtobool + +from swift.sequence_parallel import sequence_parallel +from swift.utils import get_dist_setting +from .base import BaseLoss + + +# Code borrowed from sentence_transformers +class SiameseDistanceMetric(Enum): + """The metric for the contrastive loss""" + + EUCLIDEAN = lambda x, y: F.pairwise_distance(x, y, p=2) # noqa + MANHATTAN = lambda x, y: F.pairwise_distance(x, y, p=1) # noqa + COSINE_DISTANCE = lambda x, y: 1 - F.cosine_similarity(x, y) # noqa + + +def _parse_pair_sentence(outputs): + if isinstance(outputs, dict): + last_hidden_state = outputs['last_hidden_state'] + else: + last_hidden_state = outputs + batch_size = last_hidden_state.shape[0] + shape_len = len(last_hidden_state.shape) + first_sentence = list(range(0, batch_size, 2)) + second_sentence = list(range(1, batch_size, 2)) + if shape_len == 3: + sentence1 = last_hidden_state[first_sentence][:, 0].squeeze(dim=1) + sentence2 = last_hidden_state[second_sentence][:, 0].squeeze(dim=1) + else: + sentence1 = last_hidden_state[first_sentence] + sentence2 = last_hidden_state[second_sentence] + return sentence1, sentence2 + + +class CosineSimilarityLoss(BaseLoss): + + def __call__(self, outputs, labels, **kwargs) -> torch.Tensor: + # You need to return a scalar representing the loss. + cos_score_transformation = nn.Identity() + loss_fct = MSELoss() + sentence1, sentence2 = _parse_pair_sentence(outputs) + output = cos_score_transformation(torch.cosine_similarity(sentence1, sentence2)) + return loss_fct(output, labels.to(output.dtype).view(-1)) + + +class ContrastiveLoss(BaseLoss): + + def __call__(self, outputs, labels, **kwargs) -> torch.Tensor: + sentence1, sentence2 = _parse_pair_sentence(outputs) + distance_metric = SiameseDistanceMetric.COSINE_DISTANCE + distances = distance_metric(sentence1, sentence2) + margin = 0.5 + labels = labels.to(sentence1.dtype) + losses = 0.5 * (labels * distances.pow(2) + (1 - labels) * F.relu(margin - distances).pow(2)) + return losses.mean() + + +class OnlineContrastiveLoss(BaseLoss): + + def __call__(self, outputs, labels, **kwargs) -> torch.Tensor: + sentence1, sentence2 = _parse_pair_sentence(outputs) + distance_metric = SiameseDistanceMetric.COSINE_DISTANCE + distance_matrix = distance_metric(sentence1, sentence2) + negs = distance_matrix[labels == 0] + poss = distance_matrix[labels == 1] + + # select hard positive and hard negative pairs + negative_pairs = negs[negs < (poss.max() if len(poss) > 1 else negs.mean())] + positive_pairs = poss[poss > (negs.min() if len(negs) > 1 else poss.mean())] + + positive_loss = positive_pairs.pow(2).sum() + margin = 0.5 + negative_loss = F.relu(margin - negative_pairs).pow(2).sum() + loss = positive_loss + negative_loss + return loss + + +def _parse_multi_negative_sentences(sentences, labels, hard_negatives=None): + split_indices = torch.nonzero(labels, as_tuple=False).squeeze().tolist() + if isinstance(split_indices, int): + split_indices = [split_indices] + split_indices.append(len(labels)) + split_indices = np.array(split_indices) + np.array(list(range(len(split_indices)))) + split_tensors = [] + + for i in range(len(split_indices) - 1): + start = split_indices[i] + end = split_indices[i + 1] + split_part = sentences[start:end] + if hard_negatives is not None: + negatives = len(split_part) - 2 + assert negatives > 0 + if negatives > hard_negatives: + split_part = split_part[:hard_negatives + 2] + elif negatives < hard_negatives: + selected = np.random.choice(list(range(negatives)), size=hard_negatives - negatives, replace=True) + selected += 1 # skip positive + split_part = torch.cat((split_part, split_part[selected]), dim=0) + split_tensors.append(split_part) + return split_tensors + + +class InfonceLoss(BaseLoss): + + def __call__(self, outputs, labels, **kwargs) -> torch.Tensor: + temperature = float(os.environ.get('INFONCE_TEMPERATURE', '0.1')) # temperature + # calculate CE across the batch, meaning all samples will be negative except the matching positive + use_batch = strtobool(os.environ.get('INFONCE_USE_BATCH', 'True')) + hard_negatives = os.environ.get('INFONCE_HARD_NEGATIVES', None) # how many negative prompts kept in one sample + # mask out fake negatives + infonce_mask_fake_negative = strtobool(os.environ.get('INFONCE_MASK_FAKE_NEGATIVE', 'False')) + fake_neg_margin = float(os.environ.get('INFONCE_FAKE_NEG_MARGIN', '0.1')) + # enhanced components to align with Qwen3-Embedding denominator; controlled individually + # defaults set to False for backward compatibility + infonce_include_qq = strtobool(os.environ.get('INFONCE_INCLUDE_QQ', 'False')) + infonce_include_dd = strtobool(os.environ.get('INFONCE_INCLUDE_DD', 'False')) + if hard_negatives is not None: + hard_negatives = int(hard_negatives) + if self.is_megatron: + from megatron.core import mpu + rank, world_size = mpu.get_data_parallel_rank(), mpu.get_data_parallel_world_size() + else: + rank, _, world_size, _ = get_dist_setting() + # repeat of anchor(1)+positive(1)+negatives(n) + sentences = outputs['last_hidden_state'] + + if world_size > 1 and use_batch: + if getattr(sequence_parallel, 'dp_group', None) is not None: + all_sentences = sequence_parallel._gather_object_dp(sentences.unsqueeze(0)) + labels = sequence_parallel._gather_object_dp(labels) + rank = sequence_parallel.dp_rank + elif self.is_megatron: + from megatron.core import mpu + dp_group = mpu.get_data_parallel_group() + shapes = [sentences.new_empty((2, ), dtype=torch.long) for _ in range(world_size)] + dist.all_gather( + shapes, + sentences.new_tensor(sentences.shape, dtype=torch.long), + group=dp_group, + ) + all_sentences = [sentences.new_empty(shape.tolist()) for shape in shapes] + dist.all_gather( + all_sentences, + sentences, + group=dp_group, + ) + else: + # gather all the sentences and labels across the gpus when calculate loss across all batches of all gpus + all_sentences = gather_object(sentences.unsqueeze(0)) + labels = gather_object(labels) + # override the gathered one + all_sentences[rank] = sentences + for idx in range(len(all_sentences)): + if idx == rank: + continue + # we don't calculate grad from other gpus + all_sentences[idx] = all_sentences[idx].detach().to(sentences.device) + sentences = torch.cat(all_sentences, dim=0) + labels = [tensor.to(sentences.device) for tensor in labels] + labels = torch.stack(labels, dim=0) + + # split tensors into single sample + # for example: batch_size=2 with tensor anchor(1)+positive(1)+negatives(3) + anchor(1)+positive(1)+negatives(2) + # labels will be [1,0,0,0,1,0,0], meaning 1 positive, 3 negatives, 1 positive, 2 negatives + split_tensors = _parse_multi_negative_sentences(sentences, labels, hard_negatives) + loss = 0 + can_batched = hard_negatives is not None + if hard_negatives is None and len(set([s.shape[0] for s in split_tensors])) == 1: + # all tensors have the same batch size + can_batched = True + if not use_batch: + # only calculate loss inside one sample + if can_batched: + # negative numbers are equal + # [B, neg+2, D] + sentences = torch.stack(split_tensors, dim=0) + # [B, 1, D] * [B, neg+1, D] + similarity_matrix = torch.matmul(sentences[:, 0:1], sentences[:, 1:].transpose(1, 2)) / temperature + # The positive one is the first element + labels = torch.zeros(len(split_tensors), dtype=torch.int64).to(sentences.device) + loss = nn.CrossEntropyLoss()(similarity_matrix.squeeze(1), labels) + else: + # the negative numbers may be different, use for loop + for tensor in split_tensors: + # [D] * [neg+1, D] + similarity_matrix = torch.matmul(tensor[0], tensor[1:].T) / temperature + # The positive one is the first element + labels = torch.tensor(0).to(tensor.device) + loss += nn.CrossEntropyLoss()(similarity_matrix, labels) + # avg between all batches in one gpu + loss /= len(split_tensors) + else: + if can_batched: + # [B, neg+2, D] + sentences = torch.stack(split_tensors, dim=0) + # base q->d similarities (includes own positive and all in-batch documents) + queries = sentences[:, 0].squeeze(1) # [B, D] + docs_all = sentences[:, 1:].reshape(-1, sentences.size(2)) # [B*(neg+1), D] + qd_matrix = torch.matmul(queries, docs_all.T) # [B, B*(neg+1)] + # target indices: start of each group's document block (its positive) + labels = torch.tensor(range(0, + sentences.size(0) * (sentences.size(1) - 1), + sentences.size(1) - 1)).view(-1).to(sentences.device) + + logits_list = [qd_matrix] + + if infonce_include_qq: + # q->q similarities; exclude self via -inf on diagonal to avoid accidental positives + qq_matrix = torch.matmul(queries, queries.T) # [B, B] + qq_matrix = qq_matrix.clone() + qq_matrix.fill_diagonal_(float('-inf')) + logits_list.append(qq_matrix) + + if infonce_include_dd: + # d+ -> d (doc-doc) similarities; exclude self-positive column per row + pos_docs = sentences[:, 1].squeeze(1) # [B, D] + dd_matrix = torch.matmul(pos_docs, docs_all.T) # [B, B*(neg+1)] + # mask self positive per row: column index = row_idx * (neg+1) + block = sentences.size(1) - 1 # (neg+1) + if block > 0: + row_idx = torch.arange(dd_matrix.size(0), device=dd_matrix.device) + col_idx = row_idx * block + dd_matrix[row_idx, col_idx] = float('-inf') + logits_list.append(dd_matrix) + + if infonce_mask_fake_negative: + # thresholds derived from positive q->d scores per row + row_idx = torch.arange(qd_matrix.size(0), device=qd_matrix.device) + pos_scores = qd_matrix[row_idx, labels] + thresholds = pos_scores.view(-1, 1).detach() + fake_neg_margin + + # qd block mask + qd_block = qd_matrix.clone() + qd_mask = qd_block > thresholds + qd_block[qd_mask] = float('-inf') + + components = [qd_block] + + # qq block mask (if present) + if infonce_include_qq: + qq_block = qq_matrix.clone() + qq_mask = qq_block > thresholds + qq_block[qq_mask] = float('-inf') + # diagonal already masked unconditionally at construction time + components.append(qq_block) + + # dd block (if present): self-positive column already masked unconditionally + if infonce_include_dd: + # align with Qwen3-Embedding, no threshold masking for d-d + components.append(dd_matrix) + + similarity_matrix = torch.cat(components, dim=1) + else: + # concatenate all components without masking + similarity_matrix = torch.cat(logits_list, dim=1) + # temperature scaling and CE + similarity_matrix = similarity_matrix / temperature + loss = nn.CrossEntropyLoss()(similarity_matrix, labels) + else: + all_tensors = [] + for tensor in split_tensors: + all_tensors.append(tensor[1:]) + # cat all neg+1 tensors + sentences = torch.cat(all_tensors, dim=0) + # prepare query anchors list if q-q is included + if infonce_include_qq: + queries_all = torch.stack([t[0] for t in split_tensors], dim=0) # [B, D] + length = 0 + for idx, tensor in enumerate(split_tensors): + # [D] * [B*(neg+1), D], neg numbers are different + qd_vec = torch.matmul(tensor[0], sentences.T) + target = torch.tensor(length).to(tensor.device) + logits_parts = [] + + # compute threshold from positive q->d score + threshold = (qd_vec[target].detach() + fake_neg_margin) + + # qd part with masking + if infonce_mask_fake_negative: + qd_masked = torch.where(qd_vec > threshold, torch.tensor(float('-inf'), device=qd_vec.device), + qd_vec) + else: + qd_masked = qd_vec + logits_parts.append(qd_masked) + + # qq part + if infonce_include_qq: + qq_vec = torch.matmul(tensor[0], queries_all.T) # [B] + # exclude self + qq_vec = qq_vec.clone() + qq_vec[idx] = float('-inf') + if infonce_mask_fake_negative: + qq_vec = torch.where(qq_vec > threshold, torch.tensor(float('-inf'), device=qq_vec.device), + qq_vec) + logits_parts.append(qq_vec) + + # dd part + if infonce_include_dd: + dd_vec = torch.matmul(tensor[1], sentences.T) # [B*(neg+1)] + # mask self positive column for this row only (no threshold masking for d-d) + block = split_tensors[idx].size(0) - 1 # (neg+1) for this group + dd_vec[length] = float('-inf') + logits_parts.append(dd_vec) + + logits_row = torch.cat(logits_parts, dim=-1) + logits_row = logits_row / temperature + loss += nn.CrossEntropyLoss()(logits_row.unsqueeze(0), target.unsqueeze(0)) + # next positive is neg+1 + length += tensor.size(0) - 1 + loss /= len(split_tensors) + return loss diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss/mapping.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss/mapping.py new file mode 100644 index 0000000000000000000000000000000000000000..7f1233ac5678c49312910ae5417ec3fa2b9d92e2 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss/mapping.py @@ -0,0 +1,16 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .causal_lm import CustomCrossEntropyLoss +from .embedding import ContrastiveLoss, CosineSimilarityLoss, InfonceLoss, OnlineContrastiveLoss +from .reranker import ListwiseRerankerLoss, PointwiseRerankerLoss + +loss_map = { + 'cross_entropy': CustomCrossEntropyLoss, # examples + # embedding + 'cosine_similarity': CosineSimilarityLoss, + 'contrastive': ContrastiveLoss, + 'online_contrastive': OnlineContrastiveLoss, + 'infonce': InfonceLoss, + # # reranker + 'pointwise_reranker': PointwiseRerankerLoss, + 'listwise_reranker': ListwiseRerankerLoss, +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss/reranker.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss/reranker.py new file mode 100644 index 0000000000000000000000000000000000000000..7177079e631b454bc73f0e5db2f974ed8c8a3a6a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss/reranker.py @@ -0,0 +1,109 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn import CrossEntropyLoss + +from swift.utils import get_last_valid_indices +from .base import BaseLoss + + +class PointwiseRerankerLoss(BaseLoss): + + def __call__(self, outputs, labels, **kwargs) -> torch.Tensor: + logits = outputs.logits + logits = logits.squeeze(1) + labels = labels.to(logits.dtype) + loss_fct = nn.BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + return loss + + +class ListwiseRerankerLoss(BaseLoss): + + def __call__(self, outputs, labels, **kwargs): + """ + List-wise reranker loss function. + + This loss function groups samples by query based on the pattern where each group + consists of 1 positive document followed by n negative documents. It treats the + ranking task as a classification problem within each group, using cross-entropy + loss to identify the positive document among all candidates. + + Data format expected: + - labels: [1, 0, 0, 0, 1, 0, 0, ...] where 1 indicates positive, 0 indicates negative + - Each 1 is followed by its corresponding negative documents until the next 1 + + Environment variables for configuration: + - LISTWISE_RERANKER_TEMPERATURE: Temperature for softmax (default: 1.0) + - LISTWISE_RERANKER_MIN_GROUP_SIZE: Minimum group size to include (default: 2) + + Args: + outputs: Model outputs containing logits [batch_size, 1] + labels: Binary labels (1 for positive, 0 for negative) [batch_size] + + Returns: + torch.Tensor: Cross entropy loss for ranking classification + """ + logits = outputs.logits.squeeze(-1) # [batch_size] + labels = labels.float() + + # Configuration from environment variables + temperature = float(os.environ.get('LISTWISE_RERANKER_TEMPERATURE', '1.0')) + min_group_size = int(os.environ.get('LISTWISE_RERANKER_MIN_GROUP_SIZE', '2')) + + # Find positive sample indices to determine group boundaries + positive_indices = torch.nonzero(labels == 1, as_tuple=False).squeeze(-1) + + if len(positive_indices) == 0: + # No positive samples in this batch, return zero loss + return torch.tensor(0.0, device=logits.device, requires_grad=True) + + # Ensure positive_indices is 1D + if positive_indices.dim() == 0: + positive_indices = positive_indices.unsqueeze(0) + + total_loss = 0.0 + num_groups = 0 + + for i, pos_idx in enumerate(positive_indices): + # Determine group boundaries + group_start = pos_idx.item() + + # Find the end of current group (start of next group or end of batch) + if i + 1 < len(positive_indices): + group_end = positive_indices[i + 1].item() + else: + group_end = len(labels) + + # Extract group logits and labels + group_logits = logits[group_start:group_end] # [group_size] + group_labels = labels[group_start:group_end] # [group_size] + + # Skip groups that are too small + if len(group_logits) < min_group_size: + continue + + # Verify that the first sample in the group is positive + if group_labels[0] != 1: + continue # Skip malformed groups + + # Apply temperature scaling for better training dynamics + scaled_logits = group_logits / temperature + + # The positive document is always at index 0 within the group + target = torch.tensor(0, dtype=torch.long, device=logits.device) + + # Apply cross-entropy loss: positive document should have highest score + loss_fct = CrossEntropyLoss() + group_loss = loss_fct(scaled_logits.unsqueeze(0), target.unsqueeze(0)) + + total_loss += group_loss + num_groups += 1 + + if num_groups == 0: + return torch.tensor(0.0, device=logits.device, requires_grad=True) + + # Return average loss across all groups + return total_loss / num_groups diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bd171a76345a9551896d58f8dc954839bec230e1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import ALL_BASE_STRATEGY, ConfigLossScale, LossScale +from .mapping import get_loss_scale, loss_scale_map diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/agent.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/agent.py new file mode 100644 index 0000000000000000000000000000000000000000..8851aec97e2a7daefc26eebc3b74b7affbf8a6a5 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/agent.py @@ -0,0 +1,32 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +from typing import Optional + +from .base import ConfigLossScale +from .utils import calculate_loss_scale + + +class AgentFlanLossScale(ConfigLossScale): + is_binary = False + loss_scale_config = 'agentflan.json' + + def get_loss_scale(self, context: str, *, query: Optional[str] = None, **kwargs): + if isinstance(context, str): + return calculate_loss_scale(query, context, self.loss_scale_map['response'], self.loss_scale_map['query']) + return super().get_loss_scale(context) + + +class REACTLossScale(ConfigLossScale): + loss_scale_config = 'react.json' + + +class QwenLossScale(ConfigLossScale): + loss_scale_config = 'qwen.json' + + +class HermesLossScale(ConfigLossScale): + loss_scale_config = 'hermes.json' + + +class AlphaUmiLossScale(ConfigLossScale): + loss_scale_config = 'alpha_umi.json' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/base.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/base.py new file mode 100644 index 0000000000000000000000000000000000000000..c09bdcec10b2d429909c0333fede34b7eb99731e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/base.py @@ -0,0 +1,236 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import os +from typing import List, Literal, Optional, Tuple + +from swift.template import ContextType, Messages, get_last_user_round +from .utils import calculate_loss_scale + +ALL_BASE_STRATEGY = ['default', 'last_round', 'all'] + + +class LossScale: + """Base class for loss scaling in training. + + This class provides a flexible framework for controlling loss computation weights + across different parts of the context (e.g., system prompts, user queries, assistant + responses) during model training. Different strategies can be applied to selectively + train on specific portions. + + Attributes: + is_binary (bool, optional): Indicates whether loss_scale contains only 0 and 1. + If True, loss_scale will be replaced by labels to stay compatible with + acceleration techniques such as liger_kernel. + If False, an additional 'loss_scale' key will be stored and the + corresponding loss function will be used. + base_strategy (str): Base strategy for loss computation. One of 'default', + 'last_round', or 'all'. + - 'default': Only compute loss on assistant responses + - 'last_round': Only compute loss on the last round's assistant response + - 'all': Compute loss on all parts + """ + is_binary = True + + def __init__(self, base_strategy: Literal['default', 'last_round', 'all'] = 'default'): + """Initialize the loss scale object. + + Args: + base_strategy: Base strategy for loss computation. One of 'default', + 'last_round', or 'all'. Defaults to 'default'. + + Raises: + ValueError: If the provided base_strategy is not in the allowed list. + """ + if base_strategy not in ALL_BASE_STRATEGY: + raise ValueError(f'ALL_BASE_STRATEGY: {ALL_BASE_STRATEGY}, base_strategy: {base_strategy}') + self.base_strategy = base_strategy + + def get_loss_scale(self, context: str, **kwargs) -> Tuple[List[str], List[float]]: + """Calculate loss scale for the given context. + + This is a base implementation that subclasses can override to implement + custom loss scaling logic. + + Args: + context: The input context (string). + **kwargs: Additional keyword arguments, such as query (the query of the + current round). + + Returns: + A tuple containing: + - List[str]: List of contexts, potentially split into multiple parts + - List[float]: Corresponding loss scale values, one-to-one with contexts + """ + return [context], [1.] + + def __call__(self, context_list: List[str], context_types: List[ContextType], messages: Messages, + **kwargs) -> Tuple[List[str], List[float]]: + """Process the complete conversation context and return contexts with loss scales. + + This method iterates through all context segments and determines the loss scale + for each based on the context type and base strategy. It handles special cases + such as explicitly specified loss values in messages and pre-computed loss scales. + + Args: + context_list: List of context strings or dicts, each representing a segment + of the conversation. + context_types: List of context types corresponding to each context, indicating + whether it's a system prompt, user query, assistant response, etc. + messages: Complete message list containing the conversation history. + + Returns: + A tuple containing: + - List[str]: Processed context list, potentially expanded if contexts + are split into multiple parts + - List[float]: Loss scale values corresponding one-to-one with the + returned context list + """ + res_context_list = [] + res_loss_scale = [] + i = 0 + last_user_round = get_last_user_round(messages) + for context, context_type in zip(context_list, context_types): + is_last_round = 2 * i >= last_user_round + query, loss, loss_scale = None, None, None + if context_type == ContextType.RESPONSE: + query = messages[2 * i]['content'] + # Currently, we only support applying loss/mask to the response part. + loss = messages[2 * i + 1].get('loss') + loss_scale = messages[2 * i + 1].get('loss_scale') + assert context == messages[2 * i + 1]['content'] + i += 1 + if not isinstance(context, list) or (len(context) > 0 and isinstance(context[0], int)): + context = [context] + for j in range(len(context)): + new_context, new_loss_scale = self._inner_call( + context[j], context_type, query, loss[j] if isinstance(loss, list) else loss, + loss_scale[j] if isinstance(loss_scale, list) else loss_scale, is_last_round) + res_context_list += new_context + res_loss_scale += new_loss_scale + + # The values in loss_scale_list correspond one-to-one with the values in context_list. + return res_context_list, res_loss_scale + + def _inner_call(self, context, context_type, query, loss, loss_scale, is_last_round): + if isinstance(context, dict) and 'loss_scale' in context: + new_context = [[token] for token in context['token_ids']] + loss_scale = context['loss_scale'] + else: + if isinstance(context, dict) and 'token_ids' in context: + context = context['token_ids'] + is_assistant = context_type in {ContextType.RESPONSE, ContextType.SUFFIX} + if loss or loss is None and (self.base_strategy == 'all' or + (self.base_strategy == 'default' and is_assistant) or + (self.base_strategy == 'last_round' and is_assistant and is_last_round)): + if loss_scale is None: + new_context, loss_scale = self.get_loss_scale(context, query=query) + else: + new_context, loss_scale = [context], [loss_scale] + else: + new_context, loss_scale = [context], [0.] + return new_context, loss_scale + + @property + def is_binary_loss_scale(self): + """Check if loss scale values are binary (only 0 and 1).""" + return self.is_binary + + +class ConfigLossScale(LossScale): + """Loss scale class that loads configuration from a JSON file. + + This class extends LossScale to support loading predefined loss scale mappings + from a configuration file. The mappings can specify different loss weights for + different tokens or segments based on the content. + + Attributes: + loss_scale_config (str, optional): Path to the loss scale configuration file + relative to the 'config' directory. + loss_scale_map (dict, optional): Dictionary mapping loaded from the config file, + containing predefined loss scale values for specific patterns or tokens. + """ + is_binary = None + loss_scale_config = None # path + + def __init__(self, base_strategy: Literal['default', 'last_round', 'all'] = 'default'): + """Initialize the config-based loss scale object. + + Loads the loss scale configuration from a JSON file if loss_scale_config + is specified. + + Args: + base_strategy: Base strategy for loss computation. One of 'default', + 'last_round', or 'all'. Defaults to 'default'. + """ + super().__init__(base_strategy) + self.loss_scale_map = None + if self.loss_scale_config is not None: + path = os.path.dirname(os.path.abspath(__file__)) + config_path = os.path.join(path, 'config', self.loss_scale_config) + with open(config_path, 'r', encoding='utf-8') as json_file: + self.loss_scale_map = json.load(json_file) + + @property + def is_binary_loss_scale(self): + if self.is_binary is not None: + return self.is_binary + if self.loss_scale_map is None: + return True + return all(scale in {0.0, 1.0} for lst in self.loss_scale_map.values() for scale in lst) + + def get_loss_scale(self, context: str, *, query: Optional[str] = None, **kwargs): + """Calculate loss scale using the loaded configuration. + + If context is a string, uses the configuration map to calculate loss scales + based on the query and context. Otherwise, falls back to the parent class + implementation. + + Args: + context: The input context string. + query: The user query for the current round, used to determine + appropriate loss scaling based on the configuration. + + Returns: + Tuple[List[str], List[float]]: List of context segments and their + corresponding loss scale values. + """ + if isinstance(context, str): + return calculate_loss_scale(query, context, self.loss_scale_map) + return super().get_loss_scale(context) + + +class ConcatLossScale(LossScale): + """Apply multiple loss scales sequentially. + + The output segments of each underlying loss scale are fed into the next one, + and the corresponding weights are multiplied together. This makes it possible + to compose strategies such as ``hermes+ignore_empty_think``. + """ + + is_binary = None + + def __init__(self, + loss_scales: List[LossScale], + base_strategy: Literal['default', 'last_round', 'all'] = 'default'): + super().__init__(base_strategy) + assert loss_scales, 'loss_scales must be a non-empty list' + self.loss_scales = loss_scales + + def get_loss_scale(self, context, **kwargs): + contexts = [context] + weights = [1.0] + for ls in self.loss_scales: + new_contexts: List = [] + new_weights: List[float] = [] + for c, w in zip(contexts, weights): + sub_contexts, sub_weights = ls.get_loss_scale(c, **kwargs) + for sc, sw in zip(sub_contexts, sub_weights): + new_contexts.append(sc) + new_weights.append(w * sw) + contexts = new_contexts + weights = new_weights + return contexts, weights + + @property + def is_binary_loss_scale(self): + return all(ls.is_binary_loss_scale for ls in self.loss_scales) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/agentflan.json b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/agentflan.json new file mode 100644 index 0000000000000000000000000000000000000000..2751fea02b15587835f21577221d155417d129ea --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/agentflan.json @@ -0,0 +1,22 @@ +{ + "response":{ + "Name:": [1.0, 3.0], + "Action:": [1.0, 3.0], + "ACTION:": [1.0,3.0], + "Tool:": [1.0, 3.0], + "Command": [1.0, 3.0], + "Arguments:": [1.0, 3.0], + "action input": [1.0, 3.0], + "ACTION_INPUT:":[1.0, 3.0], + "Action Input:": [1.0, 3.0], + "Thought:": [1.0, 1.0], + "Final Answer:": [1.0, 1.0], + "Observation:": [2.0, 0.0] + }, + "query":{ + "What is the tool you want to use": [3.0], + "What are the required parameter names": [3.0], + "What is the value of": [3.0], + "What are the required parameter names for this tool": [3.0] + } +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/alpha_umi.json b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/alpha_umi.json new file mode 100644 index 0000000000000000000000000000000000000000..a3d7ecac1b2e5b883b3fe36ddbc351e04efffd78 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/alpha_umi.json @@ -0,0 +1,8 @@ +{ + "Action:": [2.0, 2.0], + "Action Input:": [2.0, 2.0], + "Thought:": [1.0, 1.0], + "Final Answer:": [1.0, 1.0], + "Observation:": [2.0, 0.0], + "Next:": [2.0, 2.0] +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/hermes.json b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/hermes.json new file mode 100644 index 0000000000000000000000000000000000000000..e8bfee3fc5d6cd8aa79c99f0f9b4fcd15b623645 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/hermes.json @@ -0,0 +1,3 @@ +{ + ".+?": [2.0] +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/ignore_empty_think.json b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/ignore_empty_think.json new file mode 100644 index 0000000000000000000000000000000000000000..f7655b84d49e6538248e4aff124960bdb14320c1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/ignore_empty_think.json @@ -0,0 +1,7 @@ +{ + "^\\s*\\s*": [0.0], + "^The current thinking budget is 0, so I will directly start answering the question.\\n\\s*": [0.0], + "^\\s*": [0.0], + "^<\\|channel>thought\\n": [0.0], + "^\\s*": [0.0] +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/ignore_think_prefix.json b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/ignore_think_prefix.json new file mode 100644 index 0000000000000000000000000000000000000000..e4d19c5f0bc7dc24737fa68c15c0fa5f2dd9df2a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/ignore_think_prefix.json @@ -0,0 +1,3 @@ +{ + "^\\n?": [0.0] +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/qwen.json b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/qwen.json new file mode 100644 index 0000000000000000000000000000000000000000..731ba5340387e8a3467831877fdfb1cdd19fdc90 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/qwen.json @@ -0,0 +1,6 @@ +{ + "✿FUNCTION✿:": [2.0, 2.0], + "✿ARGS✿:": [2.0, 2.0], + "✿RETURN✿:": [1.0, 1.0], + "✿RESULT✿:": [2.0, 0.0] +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/react.json b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/react.json new file mode 100644 index 0000000000000000000000000000000000000000..006f92948e1a6de28a1825fa2ef256dc1b09de81 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/config/react.json @@ -0,0 +1,7 @@ +{ + "Action:": [2.0, 2.0], + "Action Input:": [2.0, 2.0], + "Thought:": [1.0, 1.0], + "Final Answer:": [1.0, 1.0], + "Observation:": [2.0, 0.0] +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/mapping.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/mapping.py new file mode 100644 index 0000000000000000000000000000000000000000..922cd3cb727cbbf9dfbb0305ace90de06a09c80b --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/mapping.py @@ -0,0 +1,58 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .agent import AgentFlanLossScale, AlphaUmiLossScale, HermesLossScale, QwenLossScale, REACTLossScale +from .base import ALL_BASE_STRATEGY, ConcatLossScale, LossScale +from .other import IgnoreEmptyThinkLossScale, IgnoreThinkPrefixLossScale + +# Add your loss scale here, use --loss_scale xxx to train +loss_scale_map = { + 'base': LossScale, + 'ignore_empty_think': IgnoreEmptyThinkLossScale, + 'ignore_think_prefix': IgnoreThinkPrefixLossScale, + # agent + 'react': REACTLossScale, + 'hermes': HermesLossScale, + 'qwen': QwenLossScale, + 'agentflan': AgentFlanLossScale, + 'alpha_umi': AlphaUmiLossScale, +} + + +def get_loss_scale(loss_scale: str) -> LossScale: + """Factory function to create a loss scale object from a string specification. + + The loss_scale string supports the following formats (segments separated by '+'): + 1. A strategy name alone (e.g., 'default', 'last_round', 'all') - uses base LossScale + 2. A loss scale type alone (e.g., 'hermes', 'react') - uses 'default' strategy + 3. A strategy name followed by a loss scale type (e.g., 'default+react', 'last_round+qwen') + 4. Multiple loss scale types chained together, optionally led by a base strategy + (e.g., 'hermes+ignore_empty_think', 'last_round+hermes+ignore_empty_think'). + The chained loss scales are applied sequentially: each loss scale processes the + output of the previous one and the corresponding weights are multiplied together. + + Args: + loss_scale: String specifying the loss scale configuration. + + Returns: + LossScale: An instance of the appropriate LossScale subclass. When multiple loss + scale types are specified, a ``ConcatLossScale`` wrapping them is returned. + + Examples: + >>> get_loss_scale('default') # Uses default strategy with base LossScale + >>> get_loss_scale('react') # Uses default strategy with REACTLossScale + >>> get_loss_scale('last_round+hermes') # last_round strategy with HermesLossScale + >>> get_loss_scale('last_round+hermes+ignore_empty_think') # chain hermes then ignore_empty_think + """ + parts = loss_scale.split('+') + if parts[0] in ALL_BASE_STRATEGY: + base_strategy = parts[0] + ls_names = parts[1:] or ['base'] + else: + base_strategy = 'default' + ls_names = parts + if len(ls_names) == 1: + return loss_scale_map[ls_names[0]](base_strategy) + # The base_strategy is owned by the outer ConcatLossScale; sub loss scales only + # contribute their `get_loss_scale` (which does not reference base_strategy), so + # any valid placeholder ('default') is fine here. + sub_loss_scales = [loss_scale_map[name]('default') for name in ls_names] + return ConcatLossScale(sub_loss_scales, base_strategy) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/other.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/other.py new file mode 100644 index 0000000000000000000000000000000000000000..79a7ce0929d49667ed512ee9cfd6ea9146cdcf5d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/other.py @@ -0,0 +1,10 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import ConfigLossScale + + +class IgnoreEmptyThinkLossScale(ConfigLossScale): + loss_scale_config = 'ignore_empty_think.json' + + +class IgnoreThinkPrefixLossScale(ConfigLossScale): + loss_scale_config = 'ignore_think_prefix.json' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d8e8b1a8e375a64f7e450988410386c02e28171f --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/loss_scale/utils.py @@ -0,0 +1,59 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import Dict, List, Optional, Tuple + +from swift.template import split_str_parts_by + + +def calculate_loss_scale(query: str, + response: str, + response_loss_scale_map: Dict[str, list], + query_loss_scale_map: Optional[Dict[str, list]] = None) -> Tuple[List[str], List[float]]: + """Calculate the loss scale by splitting the agent response. + + This algorithm comes from paper: https://arxiv.org/pdf/2309.00986.pdf + + Agent response format: + + ```text + Thought: you should always think about what to do + Action: the action to take, should be one of the above tools[fire_recognition, + fire_alert, call_police, call_fireman] + Action Input: the input to the action + Observation: the result of the action + ... (this Thought/Action/Action Input/Observation can be repeated zero or more times) + Thought: I now know the final answer + Final Answer: the final answer to the original input question + ``` + Returns: + A tuple of agent response parts and their weights. + """ + # query loss scale map + if query_loss_scale_map is not None: + for key in query_loss_scale_map.keys(): + if key in query: + if isinstance(query_loss_scale_map[key], (float, int)): + query_loss_scale_map[key] = [query_loss_scale_map[key]] + loss_scale_value = query_loss_scale_map[key][0] + return [response], [float(loss_scale_value)] + delimiters = [k for k, v in response_loss_scale_map.items() if len(v) == 2] + if delimiters: + agent_parts = split_str_parts_by(response, delimiters) + else: + regex_delimiters = [k for k, v in response_loss_scale_map.items() if len(v) == 1] + agent_parts = split_str_parts_by(response, regex_delimiters, regex_mode=True) + weights = [] + agent_content = [] + for c in agent_parts: + if c['key'] in response_loss_scale_map: + loss_scale = response_loss_scale_map[c['key']] + assert len(loss_scale) in {1, 2}, f'loss_scale: {loss_scale}' + if len(loss_scale) == 1: + weights += loss_scale + agent_content.append(c['content']) + else: + weights += loss_scale + agent_content += [c['key'], c['content']] + else: + weights.append(1.) + agent_content.append(c['content']) + return agent_content, weights diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bfc2bc35ab6af097acff35c47d533070c94e330c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/__init__.py @@ -0,0 +1,50 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +try: + from transformers.utils import is_torch_npu_available + + if is_torch_npu_available(): + from swift.model.npu_patch.mindspeed import prepare_mindspeed_gdn_import + prepare_mindspeed_gdn_import() + # Enable Megatron on Ascend NPU + import mindspeed.megatron_adaptor # F401 + from .init import init_megatron_env + init_megatron_env() +except Exception: + # allows lint pass. + raise + +from typing import TYPE_CHECKING + +from swift.utils.import_utils import _LazyModule + +if TYPE_CHECKING: + from .arguments import (MegatronArguments, MegatronExportArguments, MegatronPretrainArguments, + MegatronRLHFArguments, MegatronSftArguments) + from .convert import convert_hf2mcore, convert_mcore2hf + from .model import get_mcore_model + from .pipelines import megatron_export_main, megatron_pretrain_main, megatron_rlhf_main, megatron_sft_main + from .trainers import MegatronDPOTrainer, MegatronTrainer + from .utils import initialize_megatron, prepare_mcore_model +else: + _import_structure = { + 'pipelines': ['megatron_sft_main', 'megatron_pretrain_main', 'megatron_rlhf_main', 'megatron_export_main'], + 'convert': ['convert_hf2mcore', 'convert_mcore2hf'], + 'utils': ['prepare_mcore_model', 'initialize_megatron'], + 'arguments': [ + 'MegatronSftArguments', 'MegatronPretrainArguments', 'MegatronRLHFArguments', 'MegatronExportArguments', + 'MegatronArguments' + ], + 'model': ['get_mcore_model'], + 'trainers': ['MegatronTrainer', 'MegatronDPOTrainer'], + } + + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ececac99b94257b11640f16b291cc548fbc5fa3e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .export_args import MegatronExportArguments +from .megatron_args import MegatronArguments +from .pretrain_args import MegatronPretrainArguments +from .rlhf_args import MegatronRLHFArguments +from .sft_args import MegatronSftArguments diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/export_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/export_args.py new file mode 100644 index 0000000000000000000000000000000000000000..40985b261e06a3fd0688184b1b473acf7bc6ddba --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/export_args.py @@ -0,0 +1,70 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +from dataclasses import dataclass +from typing import Literal, Optional + +from swift.utils import HfConfigFactory, get_logger, to_abspath +from .megatron_args import MegatronArguments +from .megatron_base_args import MegatronBaseArguments + +logger = get_logger() + + +@dataclass +class MegatronExportArguments(MegatronBaseArguments): + to_mcore: bool = False + to_hf: bool = False + test_convert_precision: bool = False + padding_free: bool = False + attention_backend: str = 'unfused' + test_convert_dtype: str = 'float32' + recompute_granularity: Literal['selective', 'full', 'none'] = 'none' + + exist_ok: bool = False + merge_lora: Optional[bool] = None + + def _init_output_dir(self): + if self.output_dir is None: + ckpt_dir = self.ckpt_dir or f'./{self.model_suffix}' + ckpt_dir, ckpt_name = os.path.split(ckpt_dir) + if self.to_mcore: + suffix = 'mcore' + elif self.to_hf: + suffix = 'hf' + self.output_dir = os.path.join(ckpt_dir, f'{ckpt_name}-{suffix}') + + self.output_dir = to_abspath(self.output_dir) + if not self.exist_ok and os.path.exists(self.output_dir): + raise FileExistsError(f'args.output_dir: `{self.output_dir}` already exists.') + logger.info(f'args.output_dir: `{self.output_dir}`') + + def _init_megatron_args(self): + self._init_output_dir() + self.test_convert_dtype = HfConfigFactory.to_torch_dtype(self.test_convert_dtype) + extra_config = MegatronArguments.load_args_config(self.ckpt_dir) + extra_config['mcore_adapter'] = self.mcore_adapter + if self.mcore_model: + extra_config['mcore_model'] = self.mcore_model + for k, v in extra_config.items(): + setattr(self, k, v) + if self.to_hf or self.to_mcore: + self._init_convert() + if self.model_info.is_moe_model is not None and self.tensor_model_parallel_size > 1: + self.sequence_parallel = True + logger.info('Settting args.sequence_parallel: True') + if self.merge_lora is None: + self.merge_lora = self.to_hf + super()._init_megatron_args() + + def _init_convert(self): + convert_kwargs = { + 'no_save_optim': True, + 'no_save_rng': True, + 'no_load_optim': True, + 'no_load_rng': True, + 'finetune': True, + } + for k, v in convert_kwargs.items(): + setattr(self, k, v) + if self.model_info.is_moe_model: + self.moe_grouped_gemm = True diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/megatron_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/megatron_args.py new file mode 100644 index 0000000000000000000000000000000000000000..5afaacd220f6b697275b7af0b449425293dcd4ed --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/megatron_args.py @@ -0,0 +1,1125 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import megatron.core +import os +import torch +from dataclasses import dataclass, field, fields +from mcore_bridge.model import get_mcore_model_type, get_model_meta +from megatron.core import mpu +from megatron.core.transformer.enums import AttnBackend +from packaging import version +from transformers.utils import is_torch_npu_available +from transformers.utils.versions import require_version +from typing import Any, Dict, List, Literal, Optional, Tuple, Union + +from swift.arguments import ModelArguments +from swift.megatron.utils import initialize_megatron +from swift.model import get_model_info_meta +from swift.utils import get_dist_setting, get_logger, json_parse_to_dict + +mcore_016 = version.parse(megatron.core.__version__) >= version.parse('0.16.0rc0') +logger = get_logger() + + +@dataclass +class RLHFMegatronArgumentsMixin: + rlhf_type: Literal['dpo', 'kto', 'grpo', 'gkd', 'rm'] = None + loss_type: Optional[str] = None # rlhf / plugins + mcore_ref_model: Optional[str] = None + mcore_ref_adapter: Optional[str] = None + + beta: Optional[float] = None + rpo_alpha: Optional[float] = None + reference_free: bool = False + label_smoothing: float = 0. + f_divergence_type: str = 'reverse_kl' + + # kto + desirable_weight: float = 1. + undesirable_weight: float = 1. + calculate_KL: Optional[bool] = None + + # rm + center_rewards_coefficient: Optional[float] = None + + # gkd + teacher_model: Optional[str] = field(default=None) + teacher_model_type: Optional[str] = field(default=None) + teacher_model_revision: Optional[str] = field(default=None) + _teacher_use_disable_adapter: bool = False + teacher_model_server: Optional[str] = field( + default=None, + metadata={ + 'help': + 'URL of the teacher model server (e.g., http://localhost:8000). ' + 'When set, teacher logprobs are fetched via API instead of loading a local model. ' + 'Supports multi-teacher via JSON.' + }) + teacher_tag_key: str = field( + default='dataset', metadata={'help': 'Column name for multi-teacher routing. Default "dataset".'}) + gkd_logits_topk: Optional[int] = None + lmbda: float = 0.5 # On-policy probability: with prob lmbda, use student-generated responses + seq_kd: bool = False # Deprecated + offload_teacher_model: bool = False # Offload teacher model to CPU to save GPU memory + sft_alpha: float = 0.0 # Weight for SFT loss in GKD (0 = pure JSD, >0 = JSD + sft_alpha * SFT) + + # OPD-RL (On-Policy Distillation as RL): a teacher (teacher_model / teacher_model_server) + # on a GRPO run turns it into OPD-RL, injecting teacher KL as the advantage. + teacher_kl_coef: float = 1.0 + + # grpo/gkd + temperature: float = 0.9 # Temperature for sampling and loss computation + + # grpo + generation_batch_size: Optional[int] = None + steps_per_generation: Optional[int] = None + num_generations: int = 8 + num_generations_eval: Optional[int] = None + max_completion_length: int = 512 + # GSPO https://arxiv.org/abs/2507.18071 + importance_sampling_level: Literal['token', 'sequence', 'sequence_token'] = 'token' + + # SAPO https://arxiv.org/abs/2511.20347 + # Temperature parameters for soft adaptive gate + tau_pos: float = 1.0 + tau_neg: float = 1.05 + + # REAL https://arxiv.org/abs/2602.05630 + real_tau: float = 0.5 + + # FIPO https://arxiv.org/abs/2603.19835 + fipo_decay_rate: float = 32.0 + fipo_clip_range: Optional[float] = 0.2 + fipo_clip_high_only: bool = True + fipo_safety_threshold: Optional[float] = 4.0 + + epsilon: float = 0.2 + epsilon_high: Optional[float] = None + delta: Optional[float] = None + top_k: int = -1 + top_p: float = 1.0 + repetition_penalty: float = 1. + + use_vllm: bool = True + use_ray: bool = False + vllm_mode: Optional[Literal['server', 'colocate']] = None + + vllm_enable_prefix_caching: bool = True + vllm_gpu_memory_utilization: float = 0.9 + vllm_tensor_parallel_size: int = 1 + vllm_max_model_len: Optional[int] = None + vllm_enforce_eager: bool = False + vllm_limit_mm_per_prompt: Optional[Union[dict, str]] = None # '{"image": 5, "video": 2}' + vllm_disable_cascade_attn: bool = False + vllm_max_num_seqs: Optional[int] = None + vllm_mm_processor_cache_gb: Optional[float] = None + vllm_engine_kwargs: Optional[Union[dict, str]] = None + vllm_enable_lora: bool = False + + sleep_level: Literal[0, 1, 2] = 0 + offload_optimizer: bool = False + offload_model: bool = False + offload_bridge: bool = False + + vllm_server_base_url: Optional[List[str]] = None + vllm_server_host: Optional[List[str]] = None + vllm_server_port: List[int] = field(default_factory=lambda: [8000]) + vllm_server_timeout: float = 240.0 + vllm_server_group_port: Optional[List[int]] = None + + reward_funcs: List[str] = field(default_factory=list) + reward_weights: List[float] = None + # see details in swift/rewards/orm.py + # cosine reward, https://arxiv.org/abs/2502.03373 + cosine_min_len_value_wrong: float = -0.5 # r^w_0 in paper, Reward for wrong answers with zero completion length. + cosine_max_len_value_wrong: float = 0.0 # r^w_L in paper, Reward for wrong answers with max completion length. + cosine_min_len_value_correct: float = 1.0 # r^c_0 in paper, Reward for correct answers with zero completion length. + cosine_max_len_value_correct: float = 0.5 # r^c_L in paper, Reward for correct answers with max completion length. + cosine_max_len: Optional[int] = None # Lmax in paper, default equal to max_completion_length + # repetition penalty, https://arxiv.org/abs/2502.03373 + repetition_n_grams: int = 3 + repetition_max_penalty: float = -1.0 + # soft_overlong, https://arxiv.org/abs/2503.14476 + soft_max_length: Optional[int] = None + soft_cache_length: Optional[int] = None + # DAPO, https://arxiv.org/abs/2503.14476 + dynamic_sample: bool = False + max_resample_times: int = 3 + overlong_filter: bool = False + + # Dr. GRPO, https://arxiv.org/abs/2503.20783 + # GDPO: normalize each reward function separately + scale_rewards: Literal['none', 'group', 'batch', 'gdpo'] = 'group' + + # RLOO / REINFORCE++ + advantage_estimator: Literal['grpo', 'rloo', 'reinforce_plus_plus'] = 'grpo' + kl_in_reward: bool = False + + wandb_log_unique_prompts: Optional[bool] = None + log_completions: bool = False + + rollout_importance_sampling_mode: Optional[Literal['token_truncate', 'token_mask', 'sequence_truncate', + 'sequence_mask']] = None + rollout_importance_sampling_threshold: float = 2.0 + log_rollout_offpolicy_metrics: bool = False + # Off-Policy Sequence Masking: mask out sequences that deviate too much from rollout policy + # If set, compute mean(rollout_per_token_logps - per_token_logps) per sequence, + # and mask sequences where this delta > threshold AND advantage < 0 + # Falls back to old_per_token_logps if rollout_per_token_logps is not available + off_policy_sequence_mask_delta: Optional[float] = None + + # entropy + log_entropy: bool = False + # Beyond the 80/20 Rule, https://arxiv.org/abs/2506.01939 + top_entropy_quantile: float = 1.0 + + router_replay_mode: Literal['disabled', 'R2', 'R3'] = 'disabled' + + # ─────────────────────────── Not Supported Yet ─────────────────────────── + + # reward model + reward_model: Optional[List[str]] = None + reward_model_plugin: Optional[List[str]] = None + # sync ref model + sync_ref_model: bool = False + ref_model_sync_steps: int = 512 + ref_model_mixup_alpha: float = 0.6 + + async_generate: bool = False + + move_model_batches: Optional[int] = None + + # multi turn + multi_turn_scheduler: Optional[str] = None + max_turns: Optional[int] = None + completion_length_limit_scope: Literal['total', 'per_round'] = 'per_round' + vllm_server_pass_dataset: bool = False + use_gym_env: Optional[bool] = None + gym_env: Optional[str] = None + + num_iterations: int = 1 + + # dataset + dataset_shuffle: Optional[bool] = True + + def _init_kto(self): + if self.calculate_KL is None: + # Not all losses require a KL calculation + self.calculate_KL = True + if self.loss_type in ['apo_zero_unpaired']: + self.calculate_KL = False + + def __post_init__(self): + if self.rlhf_type is None: + return + default_loss_type = {'kto': 'kto', 'dpo': 'sigmoid', 'grpo': 'grpo'} + default_beta = {'gkd': 0.5, 'grpo': 0.04} + if self.beta is None: + self.beta = default_beta.get(self.rlhf_type, 0.1) + if self.loss_type is None: + self.loss_type = default_loss_type.get(self.rlhf_type) + if self.rlhf_type == 'kto': + self._init_kto() + if self.rlhf_type == 'grpo': + self._init_grpo() + if self.cosine_max_len is None: + self.cosine_max_len = self.max_completion_length + if self.vllm_limit_mm_per_prompt is not None: + self.vllm_limit_mm_per_prompt = json_parse_to_dict(self.vllm_limit_mm_per_prompt) + # Teacher setup is identical for GKD and GRPO (OPD-RL): a teacher_model / server / + # same-model LoRA self-distillation all flow through the same detection. GKD also + # allows dynamic self-distillation (no teacher at all); GRPO without a teacher is + # plain RL, so only resolve a teacher for GRPO when one is configured. + if self.rlhf_type == 'gkd' or (self.rlhf_type == 'grpo' and + (self.teacher_model is not None or self.teacher_model_server is not None)): + self._check_teacher() + if self.rlhf_type == 'grpo': + self._check_opd_rl() + if self.rlhf_type == 'gkd': + # GKD-specific: the API path only returns top-k logprobs, so gkd_logits_topk is required. + if self.teacher_model_server is not None and self.gkd_logits_topk is None: + raise ValueError('gkd_logits_topk is required when using teacher_model_server') + + # Validate gkd_logits_topk + if self.gkd_logits_topk is not None and self.gkd_logits_topk <= 0: + raise ValueError(f'gkd_logits_topk must be a positive integer, got {self.gkd_logits_topk}') + + # seq_kd (teacher-generated responses) is not implemented; raise early. + if self.seq_kd: + raise NotImplementedError('seq_kd=True (Sequential KD with teacher generation) is not supported.') + + self.num_generations = 1 + self._init_generation_batch_params() + + if self.rlhf_type in ['grpo', 'gkd']: + self.vllm_engine_kwargs = json_parse_to_dict(self.vllm_engine_kwargs) + if self.use_vllm and os.getenv('SWIFT_AUDIO_LOAD_BACKEND') is None: + # align with vLLM audio load backend + os.environ['SWIFT_AUDIO_LOAD_BACKEND'] = 'soundfile_pyav' + + @staticmethod + def resolve_generation_batch_size( + generation_batch_size: Optional[int], + steps_per_generation: Optional[int], + global_batch_size: int, + num_generations: int, + ) -> Tuple[int, int]: + """Resolve generation_batch_size and steps_per_generation from user inputs. + + Returns (generation_batch_size, steps_per_generation). + """ + if generation_batch_size is None and steps_per_generation is None: + steps_per_generation = 1 + generation_batch_size = global_batch_size * steps_per_generation + elif generation_batch_size is not None and steps_per_generation is not None: + expected = global_batch_size * steps_per_generation + if generation_batch_size != expected: + raise ValueError(f'generation_batch_size ({generation_batch_size}) must equal ' + f'global_batch_size ({global_batch_size}) * steps_per_generation ' + f'({steps_per_generation}) = {expected}.') + elif generation_batch_size is not None: + if generation_batch_size % global_batch_size != 0: + raise ValueError(f'generation_batch_size ({generation_batch_size}) ' + f'must be divisible by global_batch_size ({global_batch_size})') + steps_per_generation = generation_batch_size // global_batch_size + else: + generation_batch_size = global_batch_size * steps_per_generation + + if steps_per_generation <= 0: + raise ValueError(f'steps_per_generation must be > 0, got {steps_per_generation}.') + + if num_generations > 1 and generation_batch_size % num_generations != 0: + raise ValueError(f'generation_batch_size ({generation_batch_size}) must be divisible by ' + f'num_generations ({num_generations}).') + + return generation_batch_size, steps_per_generation + + @staticmethod + def validate_batch_dp_alignment( + generation_batch_size: int, + num_generations: int, + dp_size: int, + micro_batch_size: int, + world_size: int, + ) -> None: + """Validate that batch parameters are correctly aligned with DP parallelism. + + Reusable across both worker-side (torch.distributed) and pipeline-side + (Ray config-based dp estimation) contexts. + """ + num_rollout_prompt = generation_batch_size // num_generations + if num_rollout_prompt % dp_size != 0: + raise ValueError(f'num_rollout_prompt ({num_rollout_prompt}) = generation_batch_size ' + f'({generation_batch_size}) // num_generations ({num_generations}) ' + f'must be divisible by dp_size ({dp_size}).') + + per_device_num_rollout_prompt = num_rollout_prompt // dp_size + if per_device_num_rollout_prompt < 1: + raise ValueError(f'per_device_num_rollout_prompt ({per_device_num_rollout_prompt}) must be >= 1, ' + f'please adjust generation_batch_size/steps_per_generation/num_generations.') + + if per_device_num_rollout_prompt % micro_batch_size != 0: + raise ValueError(f'Per-device rollout prompt count ({per_device_num_rollout_prompt}) = ' + f'(generation_batch_size ({generation_batch_size}) // ' + f'num_generations ({num_generations})) // dp_size ({dp_size}) ' + f'must be divisible by micro_batch_size ({micro_batch_size}).') + + per_device_generation_batch_size = generation_batch_size // world_size + if per_device_generation_batch_size < 1: + raise ValueError(f'per_device_generation_batch_size ({per_device_generation_batch_size}) must be >= 1.') + + def _init_generation_batch_params(self): + self.generation_batch_size, self.steps_per_generation = self.resolve_generation_batch_size( + self.generation_batch_size, self.steps_per_generation, self.global_batch_size, self.num_generations) + + if torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + dp_size = world_size // ( + self.pipeline_model_parallel_size * self.tensor_model_parallel_size * self.context_parallel_size) + self.validate_batch_dp_alignment(self.generation_batch_size, self.num_generations, dp_size, + self.micro_batch_size, world_size) + self.per_device_generation_batch_size = self.generation_batch_size // world_size + + def _check_teacher(self): + """Resolve the teacher (shared by GKD and GRPO/OPD-RL). + + Detects the three teacher modes and sets ``_teacher_use_disable_adapter``: + - separate teacher_model / teacher_model_server, + - same-model LoRA self-distillation (disable_adapter, no extra model), + - dynamic self-distillation (no teacher -> teacher == current student weights). + """ + if self.teacher_model is not None and self.teacher_model_server is not None: + raise ValueError('setting both `teacher_model` and `teacher_model_server` is not supported.') + + # Fail fast: the Ray pipeline only supports a colocated teacher_model (see + # swift/ray/megatron/grpo_trainer.py), so reject teacher_model_server at parse time. + if self.use_ray and self.teacher_model_server is not None: + raise ValueError('teacher_model_server is not supported with use_ray') + + # Validate teacher_model_server: accept single URL or JSON multi-teacher config. + if self.teacher_model_server is not None: + from swift.rlhf_trainers.gkd_helpers import parse_teacher_model_server + parse_teacher_model_server(self.teacher_model_server) + + self._teacher_use_disable_adapter = False + if self.teacher_model is not None and self.teacher_model == self.model: + if self.tuner_type == 'lora': + logger.info('LoRA + same teacher_model: using disable_adapter() for fixed teacher (no extra model).') + self._teacher_use_disable_adapter = True + self.teacher_model = None + self.teacher_model_dir = None + # Full training + same teacher_model: a separate frozen copy is loaded as the fixed teacher. + + def _check_opd_rl(self): + """Fail-fast OPD-RL (teacher distillation on Megatron GRPO) parameter compatibility. + + Mirrors ``RLHFArguments._check_opd_rl``: the teacher signal is a *per-token* advantage, so + features that need a *per-sequence* advantage (sign-based) or reward variance are rejected. + Called after ``_init_grpo`` so ``loss_type`` / ``scale_rewards`` are already resolved. + """ + if self.loss_type in ['real', 'fipo']: + raise ValueError(f'OPD-RL (teacher) does not support loss_type={self.loss_type!r} ' + '(it needs a per-sequence advantage). Use grpo/bnpo/dr_grpo/dapo/cispo/sapo.') + if self.off_policy_sequence_mask_delta is not None: + raise ValueError('OPD-RL (teacher) does not support off_policy_sequence_mask_delta ' + '(it needs a per-sequence advantage).') + if not self.reward_funcs: + if self.dynamic_sample: + raise ValueError('dynamic_sample requires reward_funcs (it filters groups by reward std); ' + 'pure OPD-RL distillation has no reward variance.') + if self.scale_rewards == 'gdpo': + raise ValueError("scale_rewards='gdpo' requires reward_funcs; pure OPD-RL distillation has none.") + + def _init_grpo(self): + + def _check_not_supported(): + if self.async_generate: + raise ValueError('async_generate is not supported for Megatron GRPO right now') + if self.sync_ref_model: + raise ValueError('sync_ref_model is not supported for Megatron GRPO right now') + if self.num_iterations > 1: + raise ValueError('num_iterations > 1 is not supported for Megatron GRPO right now') + + if self.loss_type == 'real': + assert self.micro_batch_size % self.num_generations == 0, \ + (f'"REAL loss requires that the training micro_batch_size ({self.micro_batch_size}) ' + f'is a multiple of num_generations ({self.num_generations}). Please adjust your batch parameters.') + + _check_not_supported() + if self.dataset_shuffle is not None: + self.train_dataloader_shuffle = self.dataset_shuffle + self._init_generation_batch_params() + self.remove_unused_columns = False + logger.info(f'Setting args.remove_unused_columns: {self.remove_unused_columns}') + if self.truncation_strategy is None: + self.truncation_strategy = 'left' + if self.truncation_strategy not in {'left', 'delete'}: + raise ValueError("GRPO requires `truncation_strategy 'left' or 'delete'`, " + f"Current value: `truncation_strategy='{self.truncation_strategy}'`.") + # disable normalization, REAL https://arxiv.org/abs/2602.05630 + if self.loss_type == 'real': + self.scale_rewards = 'none' + logger.warning( + f"[REAL] scale_rewards='{self.scale_rewards}' is ignored. " + "It will be forced to 'none' because 'loss_type = real' does not support reward normalization.") + + if self.beta is None: + self.beta = 0.04 # https://arxiv.org/abs/2402.03300 + if self.async_generate: + logger.info('Using async mode. This is a approximate version which ' + 'will use the old weights to generate responses to accelerate. ' + 'This will ignore the `CLIP` of advantages, if you found the training ' + 'is unstable, you may consider using --async_generate false.') + if 'soft_overlong' in self.reward_funcs: + assert self.soft_cache_length is not None, \ + 'The soft_cache_length must be set when using soft overlong rewards.' + if self.soft_max_length is None: + self.soft_max_length = self.max_completion_length + logger.info(f'Auto-configured soft_max_length = max_completion_length {self.max_completion_length}') + if not self.use_ray: + assert self.use_vllm, 'use_vllm must be True for Megatron GRPO' + + # Mirror deploy_args: gym_env implies use_gym_env unless the user said otherwise. + if self.use_gym_env is None and self.gym_env is not None: + self.use_gym_env = True + + +@dataclass +class MegatronTunerMixin: + tuner_type: Literal['lora', 'full', 'lora_llm'] = 'full' + freeze_llm: bool = False + freeze_vit: bool = True + freeze_aligner: bool = True + # full + freeze_parameters: List[str] = field(default_factory=list) + freeze_parameters_regex: Optional[str] = None + freeze_parameters_ratio: float = 0. # 0 ~ 1 + trainable_parameters: List[str] = field(default_factory=list) + trainable_parameters_regex: Optional[str] = None + # lora + target_modules: List[str] = field(default_factory=lambda: ['all-linear']) + target_regex: Optional[str] = None + modules_to_save: List[str] = field(default_factory=list) + + # lora + lora_rank: int = 8 + lora_alpha: int = 32 + lora_dropout: float = 0.05 + lora_bias: Literal['none', 'all'] = 'none' + lora_dtype: Literal['float16', 'bfloat16', 'float32', None] = None + use_rslora: bool = False + + def __post_init__(self): + if 0 < self.freeze_parameters_ratio < 1 and self.pipeline_model_parallel_size > 1: + raise ValueError('`freeze_parameters_ratio` is not supported when `pipeline_model_parallel_size` > 1') + if self.target_regex: + self.target_modules = self.target_regex + + +@dataclass +class MegatronArguments(RLHFMegatronArgumentsMixin, MegatronTunerMixin): + # training + micro_batch_size: int = 1 + global_batch_size: int = 16 + recompute_granularity: Literal['selective', 'full', 'none'] = 'selective' + recompute_method: Literal['uniform', 'block'] = None + recompute_num_layers: Optional[int] = None + recompute_modules: List[str] = field(default_factory=lambda: ['core_attn']) + train_iters: Optional[int] = None + num_train_epochs: Optional[int] = None + + masked_softmax_fusion: bool = True + bias_dropout_fusion: bool = True + bias_activation_fusion: bool = True + apply_rope_fusion: bool = False + gradient_accumulation_fusion: bool = True + cross_entropy_loss_fusion: bool = True + cross_entropy_fusion_impl: Literal['native', 'te'] = 'native' + calculate_per_token_loss: Optional[bool] = None + attention_backend: str = 'flash' # flash, fused, unfused, local, auto + optimizer: Literal['adam', 'sgd', 'muon', 'dist_muon'] = 'adam' + optimizer_cpu_offload: bool = False + optimizer_offload_fraction: float = 1. + optimizer_cuda_graph: bool = False + use_precision_aware_optimizer: bool = False + main_grads_dtype: Literal['fp32', 'bf16'] = 'fp32' + main_params_dtype: Literal['fp32', 'fp16'] = 'fp32' + exp_avg_dtype: Literal['fp32', 'fp16', 'bf16', 'fp8'] = 'fp32' + exp_avg_sq_dtype: Literal['fp32', 'fp16', 'bf16', 'fp8'] = 'fp32' + manual_gc: bool = False + manual_gc_steps: int = 0 + manual_gc_eval: bool = True + + # data + seed: int = 42 + train_dataloader_shuffle: bool = True + dataloader_num_workers: int = 4 + dataloader_pin_memory: bool = True + dataloader_persistent_workers: bool = False + dataloader_prefetch_factor: int = 2 + data_sharding: bool = False + group_by_length: bool = False + te_rng_tracker: bool = False + data_parallel_random_init: Optional[bool] = False + padding_free: bool = True + mlp_padding_free: bool = False + + # learning rate + lr_warmup_init: float = 0. + lr: Optional[float] = None + lr_decay_style: Literal['constant', 'linear', 'cosine', 'inverse-square-root', 'WSD'] = 'cosine' + # The default is None, which will be set to `train_iters`. + lr_decay_iters: Optional[int] = None + lr_warmup_iters: int = 0 + lr_warmup_fraction: Optional[float] = None + min_lr: float = 0 + # wsd + lr_wsd_decay_style: Literal['exponential', 'linear', 'cosine', 'minus_sqrt'] = 'exponential' + lr_wsd_decay_iters: Optional[int] = None + + # regularization + weight_decay: float = 0.1 + weight_decay_incr_style: Literal['constant', 'linear', 'cosine'] = 'constant' + start_weight_decay: Optional[float] = None + end_weight_decay: Optional[float] = None + clip_grad: float = 1. + adam_beta1: float = 0.9 + adam_beta2: float = 0.95 + adam_eps: float = 1e-8 + sgd_momentum: float = 0.9 + muon_momentum: float = 0.9 + muon_split_qkv: bool = True + muon_use_nesterov: bool = False + muon_scale_mode: Literal['spectral', 'unit_rms_norm', 'shape_scaling'] = 'spectral' + muon_fp32_matmul_prec: Literal['low', 'medium', 'high'] = 'medium' + muon_coefficient_type: str = 'quintic' + muon_num_ns_steps: int = 5 + muon_tp_mode: Literal['blockwise', 'duplicated', 'distributed'] = 'blockwise' + muon_extra_scale_factor: float = 1. + muon_scalar_optimizer: str = 'adam' + + # checkpoint + output_dir: Optional[str] = None + save_steps: int = 500 + no_save_optim: bool = False + no_save_rng: bool = False + mcore_model: Optional[str] = None + mcore_adapter: Optional[str] = None + no_load_optim: bool = False + no_load_rng: bool = False + finetune: bool = True + perform_initialization: bool = False + use_cpu_initialization: bool = False + async_save: bool = False + save_total_limit: Optional[int] = None + metric_for_best_model: Optional[str] = None + greater_is_better: Optional[bool] = None + + use_persistent_ckpt_worker: bool = False + dist_ckpt_save_pre_mcore_014: bool = False + dist_ckpt_optim_fully_reshardable: bool = False + distrib_optim_fully_reshardable_mem_efficient: bool = False + + # dist + local_rank: Optional[int] = None # Compatible with DeepSpeed launch + ddp_timeout: int = 18000000 + ddp_backend: Literal['nccl', 'gloo'] = 'nccl' + use_distributed_optimizer: bool = True + # megatron-fsdp + use_megatron_fsdp: bool = False + data_parallel_sharding_strategy: Literal['no_shard', 'optim', 'optim_grads', + 'optim_grads_params'] = 'optim_grads_params' + tensor_model_parallel_size: int = 1 + pipeline_model_parallel_size: int = 1 + decoder_first_pipeline_num_layers: Optional[int] = None + decoder_last_pipeline_num_layers: Optional[int] = None + account_for_embedding_in_pipeline_split: bool = False + account_for_loss_in_pipeline_split: bool = False + overlap_p2p_comm: bool = True + batch_p2p_comm: Optional[bool] = None + align_param_gather: bool = True + + sequence_parallel: bool = False + context_parallel_size: int = 1 + cp_comm_type: Optional[Union[str, List[str]]] = None + cp_partition_mode: Literal['zigzag', 'contiguous'] = 'zigzag' + sequence_packing_scheduler: Optional[Literal['dp_balanced', 'default_dynamic_cp']] = None + tp_comm_overlap: bool = False + overlap_grad_reduce: bool = False + overlap_param_gather: bool = False + overlap_param_gather_with_optimizer_step: bool = False + align_grad_reduce: bool = True + # Eagerly create NCCL communicators before the training loop to avoid the lazy + # first-use allocation hitting the iteration-1 memory peak (Failed to CUDA calloc async). + nccl_comm_warmup: bool = False + virtual_pipeline_model_parallel_size: Optional[int] = None + microbatch_group_size_per_vp_stage: Optional[int] = None + pipeline_model_parallel_layout: Optional[str] = None + expert_model_parallel_size: int = 1 + expert_tensor_parallel_size: int = 1 + + # 'wandb', 'swanlab', 'tensorboard' + report_to: List[str] = field(default_factory=lambda: ['tensorboard']) + logging_steps: int = 5 + tensorboard_dir: Optional[str] = None + tensorboard_queue_size: int = 50 + wandb_project: str = 'megatron-swift' + wandb_exp_name: Optional[str] = None + swanlab_project: str = 'megatron-swift' + swanlab_exp_name: Optional[str] = None + + # evaluate + eval_iters: int = -1 + eval_steps: Optional[int] = None + + # fp8 + fp8_format: Literal['e4m3', 'hybrid'] = None + fp8_recipe: Literal['tensorwise', 'delayed', 'mxfp8', 'blockwise'] = 'delayed' + fp8_param_gather: bool = False + fp8_amax_history_len: int = 1024 + fp8_amax_compute_algo: Literal['most_recent', 'max'] = 'max' + + # fp4 + fp4_format: Literal['e2m1'] = None + fp4_recipe: Literal['nvfp4'] = 'nvfp4' + fp4_param_gather: bool = False + + # mixed precision + fp16: Optional[bool] = None + bf16: Optional[bool] = None + apply_query_key_layer_scaling: Optional[bool] = None + attention_softmax_in_fp32: bool = True + accumulate_allreduce_grads_in_fp32: bool = False + + # moe + moe_router_load_balancing_type: Optional[List[str]] = None + moe_router_dtype: Literal['none', 'fp32', 'fp64'] = 'fp32' + moe_token_dispatcher_type: Literal['allgather', 'alltoall', 'flex'] = 'alltoall' + moe_enable_deepep: bool = False + moe_grouped_gemm: bool = True + moe_permute_fusion: bool = False + moe_aux_loss_coeff: List[float] = 0. + moe_z_loss_coeff: Optional[float] = None + moe_shared_expert_overlap: bool = False + moe_layer_recompute: bool = False # compat mcore 0.12 + moe_expert_capacity_factor: Optional[float] = None + moe_pad_expert_input_to_capacity: bool = False + moe_token_drop_policy: Literal['probs', 'position'] = 'probs' + + # mtp + mtp_num_layers: Optional[int] = None + mtp_loss_scaling_factor: float = 0.1 + mtp_decoder_input_detach: bool = False + mtp_shared_weights: bool = False + + # mcore-bridge / megatron-bridge + bridge_backend: Literal['mcore-bridge', 'megatron-bridge'] = 'mcore-bridge' + model: Optional[str] = None + model_type: Optional[str] = None + save_safetensors: bool = True + save_missing_weights: bool = False + adapters: List[str] = field(default_factory=list) + ref_model: Optional[str] = None + ref_adapters: List[str] = field(default_factory=list) + use_hf: bool = False + # None: use env var `MODELSCOPE_API_TOKEN` + hub_token: Optional[str] = field( + default=None, metadata={'help': 'SDK token can be found in https://modelscope.cn/my/myaccesstoken'}) + merge_lora: Optional[bool] = None + max_shard_size: str = '5GB' + + # visual + vit_gradient_checkpointing: Optional[bool] = None + vit_gradient_checkpointing_kwargs: Optional[Union[dict, str]] = None + vit_attn_impl: Optional[str] = None + vit_lr: Optional[float] = None + aligner_lr: Optional[float] = None + + # dsa + dsa_indexer_loss_coeff: float = 0. + dsa_indexer_use_sparse_loss: bool = False + apply_dsa_kernel_fusion: bool = False + # deepseek-v4 + csa_dense_mode: bool = False + use_fused_mhc: bool = False + mhc_recompute_layer_num: Optional[int] = None + + # other + megatron_extra_kwargs: Optional[Union[dict, str]] = None + language_model_only: bool = False + check_model: bool = True + torch_dtype: Optional[Union[torch.dtype, str]] = None + rope_scaling: Optional[Union[dict, str]] = None + apply_wd_to_qk_layernorm: bool = False + linear_decoupled_in_proj: bool = False + + enable_dft_loss: bool = False + enable_channel_loss: bool = False + task_type: Literal['causal_lm', 'seq_cls', 'embedding', 'generative_reranker'] = None + num_labels: Optional[int] = None + problem_type: Literal['regression', 'single_label_classification', 'multi_label_classification'] = None + # embedding (Matryoshka Representation Learning) + # Dict[int, float], where the key is the embedding dimension and the value is the corresponding loss weight, + # e.g. '{"32": 1.0, "64": 1.0, "128": 1.0}'. + mrl_dims: Optional[Union[dict, str]] = None + save_strategy: Literal['steps', 'epoch'] = 'steps' + callbacks: List[str] = field(default_factory=list) + + @staticmethod + def load_args_config(ckpt_dir: Optional[str]) -> Dict[str, Any]: + res = {} + if ckpt_dir is None: + return res + args_path = os.path.join(ckpt_dir, 'args.json') + if os.path.exists(args_path): + with open(args_path, 'r', encoding='utf-8') as f: + old_args = json.load(f) + keys = list(f.name for f in fields(MegatronTunerMixin)) + keys += ['mcore_model', 'task_type', 'num_labels', 'bridge_backend'] + for key in keys: + old_value = old_args.get(key) + if old_value is not None: + res[key] = old_value + res.pop('mcore_adapter', None) + if res['tuner_type'] == 'full': + res.pop('mcore_model', None) + return res + + def _set_default(self): + if self.local_rank is None: + self.local_rank = get_dist_setting()[1] + if self.lr is None: + if self.tuner_type == 'full': + self.lr = 1e-5 + else: + self.lr = 1e-4 + if self.task_type is None: + self.task_type = 'causal_lm' + if self.calculate_per_token_loss is None: + self.calculate_per_token_loss = (self.task_type == 'causal_lm' and self.rlhf_type is None) + + def _init_mixed_precision(self): + ModelArguments._init_mixed_precision(self) + if self.apply_query_key_layer_scaling is None: + self.apply_query_key_layer_scaling = self.fp16 + if self.apply_query_key_layer_scaling: + os.environ['NVTE_APPLY_QK_LAYER_SCALING'] = '1' + + def _check_mcore_bridge(self): + if self.language_model_only: + if self.tuner_type == 'lora_llm': + raise ValueError('`tuner_type="lora_llm"` is not supported when `language_model_only=True`. ' + 'Please use `tuner_type="lora"` instead.') + + def _check_bridge_backend(self): + """Validate bridge_backend and associated constraints.""" + if self.bridge_backend == 'megatron-bridge': + try: + import megatron.bridge + except ImportError: + raise ImportError('bridge_backend="megatron-bridge" requires the `megatron-bridge` package. ' + 'Install it via `pip install megatron-bridge` or use bridge_backend="mcore-bridge".') + if self.tuner_type != 'full': + raise ValueError('LoRA training is not yet supported with bridge_backend="megatron-bridge". ' + 'Please use bridge_backend="mcore-bridge" for LoRA, or set tuner_type="full".') + else: + require_version('mcore-bridge>=1.5.0', 'Please install mcore-bridge via `pip install mcore-bridge -U`') + from swift.megatron.init import _patch_mcore_bridge + _patch_mcore_bridge() + self._check_mcore_bridge() + + def __post_init__(self): + if self.tuner_type != 'full': + require_version('peft>=0.15', 'Please install peft>=0.15 to use LoRA in Megatron-SWIFT.') + RLHFMegatronArgumentsMixin.__post_init__(self) + MegatronTunerMixin.__post_init__(self) + os.environ.setdefault('CUDA_DEVICE_MAX_CONNECTIONS', '1') + self._check_bridge_backend() + if self.recompute_granularity == 'none': + self.recompute_granularity = None + if self.recompute_granularity == 'selective' and self.recompute_method is not None: + raise ValueError('recompute method is not yet supported for selective recomputing granularity') + + if self.group_by_length and self.padding_free: + raise ValueError('group_by_length is not compatible with padding_free.') + self._set_default() + self._init_vpp_size() + if self.vit_gradient_checkpointing is None: + self.vit_gradient_checkpointing = not self.freeze_vit + if isinstance(self.report_to, str): + self.report_to = [self.report_to] + self.model_info, self.model_meta = get_model_info_meta( + self.model, model_type=self.model_type, use_hf=self.use_hf, hub_token=self.hub_token) + + # Megatron has a model_type parameter with the same name, so we need to avoid conflicts. + self.model_type = self.model_info.model_type + self.model_dir = self.model_info.model_dir + self.is_multimodal = self.model_meta.is_multimodal + if self.bridge_backend == 'megatron-bridge': + self.megatron_model_meta = None + if self.is_multimodal: + raise ValueError('Multimodal training is not yet supported with bridge_backend="megatron-bridge". ' + 'Please use bridge_backend="mcore-bridge" for multimodal models.') + if self.task_type not in (None, 'causal_lm'): + raise ValueError(f'task_type={self.task_type!r} is not yet supported with ' + f'bridge_backend="megatron-bridge".') + else: + self.megatron_model_meta = get_model_meta(self._get_mcore_model_type(self.model_meta)) + if self.megatron_model_meta is None: + raise ValueError(f'Model: {self.model} is not supported.') + self._init_teacher_model() + if self.apply_wd_to_qk_layernorm and self.model_type not in {'qwen3_next', 'qwen3_5', 'qwen3_5_moe'}: + raise ValueError('apply_wd_to_qk_layernorm is only supported for qwen3_next, qwen3_5 and qwen3_5_moe') + if self.pipeline_model_parallel_size == 1 and (self.decoder_first_pipeline_num_layers is not None + or self.decoder_last_pipeline_num_layers is not None): + raise ValueError('pipeline_model_parallel_size must be greater than 1 if you want to set ' + 'decoder_first_pipeline_num_layers or decoder_last_pipeline_num_layers.') + # compat megatron-core + self.fp8 = self.fp8_format + self.fp4 = self.fp4_format + + if self.megatron_extra_kwargs is not None: + self.megatron_extra_kwargs = json_parse_to_dict(self.megatron_extra_kwargs) + if self.mrl_dims is not None: + self.mrl_dims = json_parse_to_dict(self.mrl_dims) + self.mrl_dims = {int(k): float(v) for k, v in self.mrl_dims.items()} + if self.task_type not in {'causal_lm', 'generative_reranker'}: + self.untie_embeddings_and_output_weights = True + if self.vit_gradient_checkpointing_kwargs is not None: + self.vit_gradient_checkpointing_kwargs = json_parse_to_dict(self.vit_gradient_checkpointing_kwargs) + if self.gradient_accumulation_fusion: + try: + import apex + except ImportError: + logger.warning('apex is not installed, so gradient accumulation fusion is disabled.') + self.gradient_accumulation_fusion = False + self.callbacks += ['print', 'default_flow'] + self.callbacks += self.report_to + if self.save_total_limit is not None: + if self.async_save: + raise ValueError('async_save is not supported with save_total_limit.') + if self.save_total_limit < 2: + raise ValueError('save_total_limit must be greater than or equal to 2.') + if self.metric_for_best_model is None: + self.metric_for_best_model = 'reward' if self.rlhf_type == 'grpo' else 'loss' + if self.greater_is_better is None and self.metric_for_best_model is not None: + self.greater_is_better = 'loss' not in self.metric_for_best_model + if isinstance(self.ref_adapters, str): + self.ref_adapters = [self.ref_adapters] + if self.eval_steps is None: + self.eval_steps = self.save_steps + if self.merge_lora is None: + self.merge_lora = self.save_safetensors + if self.tuner_type == 'lora_llm': + if not self.is_multimodal: + raise ValueError('`tuner_type="lora_llm"` is only supported for multimodal models.') + if not self.merge_lora: + raise ValueError('`merge_lora` must be True when using `--tuner_type lora_llm`') + if not self.no_save_optim: + raise ValueError('`no_save_optim` must be True when using `--tuner_type lora_llm`') + if self.adapters or self.ref_adapters or self.mcore_adapter or self.mcore_ref_adapter: + if self.tuner_type == 'full': + self.tuner_type = 'lora' + logger.info('Setting args.tuner_type: lora') + if self.adapters: + self._load_adapter_config() + self._init_mixed_precision() + self._init_multimodal_full() + self._map_dtype() + self._init_weigh_decay() + self._init_attention_backend() + if self.sequence_parallel and self.tensor_model_parallel_size <= 1: + self.sequence_parallel = False + if isinstance(self.moe_aux_loss_coeff, list) and len(self.moe_aux_loss_coeff) == 1: + self.moe_aux_loss_coeff = self.moe_aux_loss_coeff[0] + if isinstance(self.moe_router_load_balancing_type, list) and len(self.moe_router_load_balancing_type) == 1: + self.moe_router_load_balancing_type = self.moe_router_load_balancing_type[0] + if self.tp_comm_overlap and not self.sequence_parallel: + raise ValueError('Tensor parallel communication/GEMM overlap can happen only when ' + 'sequence parallelism is enabled') + + self._check_megatron_fsdp() + self._init_distributed() + self._check_muon() + + def _check_megatron_fsdp(self): + if not self.use_megatron_fsdp: + return + # Megatron-FSDP is only compatible with the distributed optimizer. + if not self.use_distributed_optimizer: + logger.info('Megatron-FSDP is only compatible with use_distributed_optimizer=True; setting it to True.') + self.use_distributed_optimizer = True + # Only sgd/adam are supported by Megatron-FSDP. + if self.optimizer not in ('sgd', 'adam'): + raise ValueError(f'Megatron-FSDP does not support the "{self.optimizer}" optimizer yet.') + # FSDP requires CUDA_DEVICE_MAX_CONNECTIONS > 1 (or unset). SWIFT sets it to '1' by default, + # so override it here before distributed initialization. + if os.environ.get('CUDA_DEVICE_MAX_CONNECTIONS') == '1': + raise ValueError('Megatron-FSDP requires `CUDA_DEVICE_MAX_CONNECTIONS > 1`, ' + 'for example you can set `CUDA_DEVICE_MAX_CONNECTIONS=32`') + + def _init_attention_backend(self): + if self.attention_backend.startswith('flash_'): + from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils as fa_utils + + fa_version = int(self.attention_backend[len('flash_'):]) + assert fa_version in (2, 3, 4), (f'Unsupported flash attention version: {fa_version}. ' + f'Supported: flash_2, flash_3, flash_4.') + available = {2: fa_utils.is_installed, 3: fa_utils.v3_is_installed, 4: fa_utils.v4_is_installed} + if not available[fa_version]: + raise ValueError(f'flash-attn v{fa_version} is not installed. ' + f'Detected installations: FA2={available[2]}, FA3={available[3]}, FA4={available[4]}.') + + if fa_version != 2: + fa_utils.is_installed = False + if fa_version != 3: + fa_utils.v3_is_installed = False + if fa_version != 4: + fa_utils.v4_is_installed = False + logger.info(f'Forcing Flash Attention v{fa_version} as the attention backend.') + self.attention_backend = 'flash' + self.attention_backend = AttnBackend[self.attention_backend] + + def _init_distributed(self): + initialize_megatron(self) + total_model_size = ( + self.tensor_model_parallel_size * self.pipeline_model_parallel_size * self.context_parallel_size) + # world_size is initialized in initialize_megatron + self.data_parallel_size = self.world_size // total_model_size + # Gradient Accumulation + micro_batch_times_data_parallel_size = self.micro_batch_size * self.data_parallel_size + if self.global_batch_size % micro_batch_times_data_parallel_size != 0: + raise ValueError(f'global batch size ({self.global_batch_size}) is not divisible by micro batch size ' + f'({self.micro_batch_size}) times data parallel size ({self.data_parallel_size}).') + self.num_microbatches = self.global_batch_size // micro_batch_times_data_parallel_size + if self.num_microbatches == 0: + raise ValueError('global_batch_size must be >= `data_parallel_size * micro_batch_size` ' + f'to have at least one micro-batch. global_batch_size: {self.global_batch_size}, ' + f'data_parallel_size: {self.data_parallel_size}, ' + f'micro_batch_size: {self.micro_batch_size}.') + + def _get_mcore_model_type(self, model_meta): + model_type = model_meta.model_type + mcore_model_type = self.model_meta.mcore_model_type + if mcore_model_type is None: + mcore_model_type = get_mcore_model_type(model_type) + return mcore_model_type + + def _check_muon(self): + # Code borrowed from NVIDIA/Megatron-LM + if 'muon' in self.optimizer: + if not mcore_016: + raise ValueError('Muon optimizer requires "megatron-core>=0.16".') + + if self.optimizer == 'muon': + assert not self.overlap_grad_reduce, ( + 'Muon optimizer does not support overlap grad reduce. Use dist_muon instead.') + assert not self.overlap_param_gather, ( + 'Muon optimizer does not support overlap param gather. Use dist_muon instead.') + # Muon optimizer does not support distributed optimizer for now. + self.use_distributed_optimizer = False + # compat mcore 0.17 + self.muon_nesterov = self.muon_use_nesterov + + def _init_teacher_model(self): + if self.teacher_model is None: + return + self.teacher_model_info, self.teacher_model_meta = get_model_info_meta( + self.teacher_model, model_type=self.teacher_model_type, use_hf=self.use_hf, hub_token=self.hub_token) + self.teacher_model_type = self.teacher_model_info.model_type + self.teacher_model_dir = self.teacher_model_info.model_dir + if self.bridge_backend == 'megatron-bridge': + self.teacher_megatron_model_meta = None + else: + self.teacher_megatron_model_meta = get_model_meta(self._get_mcore_model_type(self.teacher_model_meta)) + if self.teacher_megatron_model_meta is None: + raise ValueError(f'Model: {self.teacher_model} is not supported.') + + def _init_vpp_size(self): + if self.pipeline_model_parallel_layout is not None: + from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout + + # Parse the input flattened layout to a list and get the vpp size. + # We will validate the layout more carefully in the TransformerConfig constructor. + num_stages = PipelineParallelLayerLayout.get_num_stages_from_str(self.pipeline_model_parallel_layout) + assert num_stages % self.pipeline_model_parallel_size == 0, ( + f'The length of pipeline_model_parallel_layout must be divisible' + f' by pipeline_model_parallel_size ({num_stages=},' + f' {self.pipeline_model_parallel_size=})') + self.virtual_pipeline_model_parallel_size = num_stages // self.pipeline_model_parallel_size + elif self.virtual_pipeline_model_parallel_size is not None: + self.virtual_pipeline_model_parallel_size = self.virtual_pipeline_model_parallel_size + if self.virtual_pipeline_model_parallel_size == 1: + self.virtual_pipeline_model_parallel_size = None + if self.virtual_pipeline_model_parallel_size is None: + self.overlap_p2p_comm = False + self.align_param_gather = False + if self.batch_p2p_comm is None: + self.batch_p2p_comm = not self.overlap_p2p_comm + + def _load_adapter_config(self): + assert len(self.adapters) == 1, 'Currently only support one adapter' + adapter_path = self.adapters[0] + adapter_config_path = os.path.join(adapter_path, 'adapter_config.json') + adapter_config = {} + if os.path.exists(adapter_config_path): + with open(adapter_config_path, 'r', encoding='utf-8') as f: + adapter_config = json.load(f) + mapping = {'r': 'lora_rank', 'bias': 'lora_bias'} + for k in ['lora_alpha', 'lora_dropout', 'use_rslora']: + mapping[k] = k + for k, v in adapter_config.items(): + if k not in mapping: + continue + k = mapping[k] + if v != getattr(self, k): + setattr(self, k, v) + logger.info(f'Setting {k}: {v}') + + def init_iters(self, train_dataset, val_dataset): + data_parallel_size = mpu.get_data_parallel_world_size() + step_batch_size = self.micro_batch_size * data_parallel_size + num_generations = self.num_generations if self.rlhf_type == 'grpo' else 1 + # TODO: Check if it causes duplicate saving at the end. + if self.save_strategy == 'epoch': + if hasattr(train_dataset, '__len__'): + dataset_sample = len(train_dataset) // step_batch_size * step_batch_size * num_generations + self.save_steps = dataset_sample // self.global_batch_size + self.eval_steps = self.save_steps + else: + raise ValueError('streaming dataset is not supported with `--save_strategy epoch`.') + if self.num_train_epochs is not None: + if hasattr(train_dataset, '__len__'): + dataset_sample = len(train_dataset) // step_batch_size * step_batch_size * num_generations + self.train_iters = dataset_sample * self.num_train_epochs // self.global_batch_size + elif self.train_iters is None: + raise ValueError( + 'You are using a streaming training dataset. Please explicitly specify `--train_iters`.') + if self.eval_iters < 0: + if val_dataset is None: + self.eval_iters = 0 + elif hasattr(val_dataset, '__len__'): + dataset_sample = len(val_dataset) // step_batch_size * step_batch_size + dataset_sample = dataset_sample * num_generations + self.eval_iters = max(dataset_sample // self.global_batch_size, 1) + else: + raise ValueError( + 'You are using a streaming validation dataset. Please explicitly specify `--eval_iters`.') + logger.info(f'Setting args.eval_iters: {self.eval_iters}') + + data_parallel_size = mpu.get_data_parallel_world_size() + step_batch_size = self.micro_batch_size * data_parallel_size + # To avoid errors caused by the validation set being insufficient to complete a single step. + if val_dataset is not None and hasattr(val_dataset, '__len__') and len(val_dataset) < step_batch_size: + val_dataset = None + if val_dataset is None: + self.eval_iters = 0 + + def _init_multimodal_full(self): + if not self.is_multimodal: + return + visual_cls = self.megatron_model_meta.visual_cls + if self.tuner_type == 'full' and self.is_multimodal and visual_cls is not None and not self.language_model_only: + vision_tower = [f'visual.{vit}' for vit in getattr(visual_cls, '_vision_tower', [])] + aligner = [f'visual.{aligner}' for aligner in getattr(visual_cls, '_aligner', [])] + generator = [f'visual.{generator}' for generator in getattr(visual_cls, '_generator', [])] + if self.freeze_llm: + self.freeze_parameters.append('language_model') + if self.freeze_vit: + self.freeze_parameters += vision_tower + if self.freeze_aligner: + self.freeze_parameters += aligner + else: + self.trainable_parameters += aligner + self.freeze_parameters += generator + if self.freeze_parameters: + logger.info(f'freeze_parameters: {self.freeze_parameters}') + if self.trainable_parameters: + logger.info(f'additional trainable_parameters: {self.trainable_parameters}') + + def _map_dtype(self): + dtype_map = { + 'fp32': torch.float32, + 'bf16': torch.bfloat16, + 'fp16': torch.float16, + 'fp8': torch.uint8, + } + self.main_grads_dtype = dtype_map[self.main_grads_dtype] + self.main_params_dtype = dtype_map[self.main_params_dtype] + self.exp_avg_dtype = dtype_map[self.exp_avg_dtype] + self.exp_avg_sq_dtype = dtype_map[self.exp_avg_sq_dtype] + if self.fp16: + self.torch_dtype = torch.float16 + elif self.bf16: + self.torch_dtype = torch.bfloat16 + if self.main_grads_dtype == torch.float32: + self.accumulate_allreduce_grads_in_fp32 = True + + def _init_weigh_decay(self): + if self.weight_decay_incr_style == 'constant': + assert self.start_weight_decay is None + assert self.end_weight_decay is None + self.start_weight_decay = self.end_weight_decay = self.weight_decay + else: + assert self.start_weight_decay is not None + assert self.end_weight_decay is not None diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/megatron_base_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/megatron_base_args.py new file mode 100644 index 0000000000000000000000000000000000000000..0f3a9c0f550eebe54d87458d8fd51b21e4357948 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/megatron_base_args.py @@ -0,0 +1,47 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +from dataclasses import dataclass + +from swift.arguments import BaseArguments +from swift.utils import get_logger +from .megatron_args import MegatronArguments, RLHFMegatronArgumentsMixin + +logger = get_logger() + + +@dataclass +class MegatronBaseArguments(MegatronArguments, BaseArguments): + + # true for ray pipeline to skip distributed init + skip_megatron_init: bool = False + + def _init_distributed(self): + if self.skip_megatron_init: + return + super()._init_distributed() + + def _init_output_dir(self): + if self.skip_megatron_init: + if self.output_dir is None: + self.output_dir = f'megatron_output/{self.model_suffix}' + os.makedirs(self.output_dir, exist_ok=True) + return + super()._init_output_dir() + + def _init_megatron_args(self): + MegatronArguments.__post_init__(self) + + def __post_init__(self): + self.sequence_parallel_size = self.context_parallel_size + if self.packing: + self.padding_free = True + BaseArguments.__post_init__(self) + self.seq_length = self.packing_length or self.max_length + if self.skip_megatron_init: + RLHFMegatronArgumentsMixin.__post_init__(self) + else: + self._init_megatron_args() + if self.streaming: + if self.dataloader_num_workers > 1: + self.dataloader_num_workers = 1 + logger.info('Using streaming dataset, setting args.dataloader_num_workers to 1.') diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/pretrain_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/pretrain_args.py new file mode 100644 index 0000000000000000000000000000000000000000..ba7a2ee229b92e83b833fa073361d83a9164e4e4 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/pretrain_args.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + +from .sft_args import MegatronSftArguments + + +@dataclass +class MegatronPretrainArguments(MegatronSftArguments): + use_chat_template: bool = False + loss_scale: str = 'all' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/rlhf_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/rlhf_args.py new file mode 100644 index 0000000000000000000000000000000000000000..22ebbc146abc0d712b3e32dcdc84731baa0aaf9a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/rlhf_args.py @@ -0,0 +1,29 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import dataclass +from typing import Literal, Optional + +from .sft_args import MegatronSftArguments + + +@dataclass +class MegatronRLHFArguments(MegatronSftArguments): + rlhf_type: Literal['dpo', 'kto', 'grpo', 'gkd', 'rm'] = 'dpo' + loss_scale: str = 'last_round' + truncation_strategy: Optional[Literal['delete', 'left', 'right', 'split', None]] = None + + calculate_per_token_loss: bool = False + + def __post_init__(self): + if self.rlhf_type == 'rm': + self.task_type = 'seq_cls' + self.num_labels = 1 + self._init_truncation_strategy() + super().__post_init__() + + def _init_truncation_strategy(self): + if self.truncation_strategy is not None: + return + if self.rlhf_type == 'grpo': + self.truncation_strategy = 'left' + else: + self.truncation_strategy = 'delete' diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/sft_args.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/sft_args.py new file mode 100644 index 0000000000000000000000000000000000000000..08fef33a5c38dc6086e1f28a98e09dcaff37634c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/arguments/sft_args.py @@ -0,0 +1,53 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import os +from dataclasses import dataclass + +from swift.utils import add_version_to_work_dir, get_logger, init_process_group, is_last_rank, to_abspath +from .megatron_base_args import MegatronBaseArguments + +logger = get_logger() + + +@dataclass +class MegatronSftArguments(MegatronBaseArguments): + add_version: bool = True + create_checkpoint_symlink: bool = False + load_args: bool = False + + def _init_output_dir(self): + init_process_group(backend=self.ddp_backend, timeout=self.ddp_timeout) + if self.output_dir is None: + self.output_dir = f'megatron_output/{self.model_suffix}' + self.output_dir = to_abspath(self.output_dir) + if self.add_version: + self.output_dir = add_version_to_work_dir(self.output_dir) + logger.info(f'args.output_dir: {self.output_dir}') + os.makedirs(self.output_dir, exist_ok=True) + + def _init_ckpt_dir(self, adapters=None): + super()._init_ckpt_dir(adapters) + if self.ckpt_dir and self.model is None: + args_path = os.path.join(self.ckpt_dir, 'args.json') + if not os.path.exists(args_path): + return + with open(args_path, 'r', encoding='utf-8') as f: + old_args = json.load(f) + self.model = old_args.get('model') + + def _init_megatron_args(self): + self._init_output_dir() + super()._init_megatron_args() + + def __post_init__(self): + self.mcore_model = to_abspath(self.mcore_model, check_path_exist=True) + super().__post_init__() + if len(self.dataset) == 0 and len(self.cached_dataset) == 0: + raise ValueError(f'self.dataset: {self.dataset}, self.cached_dataset: {self.cached_dataset}. ' + 'Please input the training dataset.') + if self.tensorboard_dir is None and self.output_dir is not None: + self.tensorboard_dir = f'{self.output_dir}/runs' + self.tensorboard_dir = to_abspath(self.tensorboard_dir) + if self.mcore_model is None and self.model is None and not self.perform_initialization: + raise ValueError('You did not pass `--mcore_model/--model` to read weights, so you need to set ' + '`--perform_initialization true` to allow the model to initialize weights properly.') diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b045ca748dc0c0de685f879c19b60aebff490f30 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/__init__.py @@ -0,0 +1,3 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .base import MegatronCallback +from .mapping import megatron_callbacks_map diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/base.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/base.py new file mode 100644 index 0000000000000000000000000000000000000000..c3dba426d3721df57a5d06a7d03ad9bb48283da6 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/base.py @@ -0,0 +1,46 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from swift.megatron.trainers import BaseMegatronTrainer + + +class MegatronCallback: + + def __init__(self, trainer: 'BaseMegatronTrainer'): + self.trainer = trainer + self.args = trainer.args + self.state = trainer.state + + def on_train_begin(self): + pass + + def on_train_end(self): + pass + + def on_step_begin(self): + pass + + def on_step_end(self): + pass + + def on_log(self, logs): + pass + + def on_eval_begin(self): + pass + + def on_eval_end(self): + pass + + def on_eval_step(self): + pass + + def on_save(self, output_dir): + """Called after save_checkpoint() returns. + + Note: When async_save is enabled, the checkpoint may not be fully + written to disk yet. Use only for non-I/O-dependent logic, or + ensure async_save is disabled if you need to read the checkpoint. + """ + pass diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/default_flow.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/default_flow.py new file mode 100644 index 0000000000000000000000000000000000000000..5253d8246753ef03e105d19074b13c21ac5ae8f3 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/default_flow.py @@ -0,0 +1,43 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import gc + +from .base import MegatronCallback + + +class DefaultFlowCallback(MegatronCallback): + + def on_train_begin(self): + args = self.args + if args.manual_gc: + gc.disable() + gc.collect() + + def on_step_end(self): + args = self.args + state = self.state + + state.consumed_train_samples += args.global_batch_size + + if state.iteration == 1 or state.iteration % args.logging_steps == 0: + state.should_log = True + if args.eval_steps and state.iteration % args.eval_steps == 0 and args.eval_iters > 0: + state.should_eval = True + if args.save_steps and state.iteration % args.save_steps == 0: + state.should_save = True + + if state.iteration >= args.train_iters: + if args.eval_iters > 0: + state.should_eval = True + state.should_save = True + if args.manual_gc and args.manual_gc_steps != 0 and state.iteration % args.manual_gc_steps == 0: + gc.collect() + + def on_eval_begin(self): + args = self.args + if args.manual_gc and args.manual_gc_eval: + gc.collect() + + def on_eval_end(self): + args = self.args + if args.manual_gc and args.manual_gc_eval: + gc.collect(generation=0) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/mapping.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/mapping.py new file mode 100644 index 0000000000000000000000000000000000000000..68e13269c173fbc5e9baab80f34374132a24aa31 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/mapping.py @@ -0,0 +1,14 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from .default_flow import DefaultFlowCallback +from .print import PrintCallback +from .swanlab import SwanlabCallback +from .tensorboard import TensorboardCallback +from .wandb import WandbCallback + +megatron_callbacks_map = { + 'print': PrintCallback, + 'default_flow': DefaultFlowCallback, + 'swanlab': SwanlabCallback, + 'wandb': WandbCallback, + 'tensorboard': TensorboardCallback, +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/print.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/print.py new file mode 100644 index 0000000000000000000000000000000000000000..01a41b385d27205326fef5b19a2ec8a11e68c119 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/print.py @@ -0,0 +1,69 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import time +import torch +from tqdm import tqdm + +from swift.megatron.utils import reduce_max_stat_across_model_parallel_group +from swift.utils import JsonlWriter, format_time, get_logger, is_last_rank +from .base import MegatronCallback + +logger = get_logger() + + +class PrintCallback(MegatronCallback): + + def __init__(self, trainer): + super().__init__(trainer) + self.training_bar = None + self.eval_bar = None + self.jsonl_writer = None + self.is_write_rank = is_last_rank() + + def on_train_begin(self): + self.training_bar = tqdm( + total=self.args.train_iters, dynamic_ncols=True, disable=not self.is_write_rank, desc='Train: ') + self.start_step = self.state.iteration + self.training_bar.update(self.state.iteration) + self.current_step = self.state.iteration + self.start_time = time.time() + logging_path = os.path.join(self.args.output_dir, 'logging.jsonl') + logger.info(f'logging_path: {logging_path}') + self.jsonl_writer = JsonlWriter(logging_path, enable_async=True, write_on_rank='last') + + def on_train_end(self): + self.training_bar.close() + self.training_bar = None + + def on_step_end(self): + n_step = self.state.iteration - self.current_step + self.current_step = self.state.iteration + self.training_bar.update(n_step) + + def on_eval_begin(self): + self.eval_bar = tqdm( + total=self.args.eval_iters, dynamic_ncols=True, disable=not self.is_write_rank, desc='Evaluate: ') + + def on_eval_end(self): + self.eval_bar.close() + self.eval_bar = None + + def on_eval_step(self): + self.eval_bar.update() + + def on_log(self, logs): + state = self.state + args = self.args + logs['iteration'] = f'{state.iteration}/{args.train_iters}' + elapsed = time.time() - self.start_time + logs['elapsed_time'] = format_time(elapsed) + n_steps = state.iteration - self.start_step + train_speed = elapsed / n_steps if n_steps > 0 else 0.0 + logs['remaining_time'] = format_time((args.train_iters - state.iteration) * train_speed) + memory = reduce_max_stat_across_model_parallel_group(torch.cuda.max_memory_reserved() / 1024**3) + logs['memory(GiB)'] = round(memory, 2) + logs['train_speed(s/it)'] = round(train_speed, 6) + logs = {k: round(v, 8) if isinstance(v, float) else v for k, v in logs.items()} + self.jsonl_writer.append(logs) + if self.is_write_rank: + self.training_bar.write(str(logs)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/swanlab.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/swanlab.py new file mode 100644 index 0000000000000000000000000000000000000000..c3f964374184b76d0dd5c11599db31c688877631 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/swanlab.py @@ -0,0 +1,35 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os + +from swift.utils import check_json_format, is_last_rank +from .base import MegatronCallback +from .utils import rewrite_logs + + +class SwanlabCallback(MegatronCallback): + + def __init__(self, trainer): + super().__init__(trainer) + args = self.args + self.config = check_json_format(vars(args)) + if args.swanlab_exp_name is None: + args.swanlab_exp_name = args.output_dir + self.save_dir = os.path.join(args.output_dir, 'swanlab') + self.writer = None + self.setup() + + def setup(self): + import swanlab + args = self.args + if is_last_rank(): + swanlab.init( + logdir=self.save_dir, + experiment_name=args.swanlab_exp_name, + project=args.swanlab_project, + config=self.config) + self.writer = swanlab + + def on_log(self, logs): + logs = rewrite_logs(logs) + if is_last_rank(): + self.writer.log(logs, step=self.state.iteration) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/tensorboard.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/tensorboard.py new file mode 100644 index 0000000000000000000000000000000000000000..9b46960d602713039d4538ccd4127257f6f307d2 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/tensorboard.py @@ -0,0 +1,32 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from swift.utils import check_json_format, is_last_rank +from .base import MegatronCallback +from .utils import rewrite_logs + + +class TensorboardCallback(MegatronCallback): + + def __init__(self, trainer): + super().__init__(trainer) + args = self.args + self.config = check_json_format(vars(args)) + self.save_dir = args.tensorboard_dir + if self.save_dir is None: + self.save_dir = f'{args.output_dir}/runs' + from torch.utils.tensorboard import SummaryWriter + self.writer = None + if is_last_rank(): + self.writer = SummaryWriter(log_dir=self.save_dir, max_queue=args.tensorboard_queue_size) + for k, v in self.config.items(): + self.writer.add_text(k, str(v), global_step=self.state.iteration) + + def on_log(self, logs): + logs = rewrite_logs(logs) + if self.writer: + for k, v in logs.items(): + self.writer.add_scalar(k, v, self.state.iteration) + + def on_train_end(self): + if self.writer: + self.writer.close() + self.writer = None diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f6e51b559166f3fc954eb1a3a846fb8cfb1a2711 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/utils.py @@ -0,0 +1,19 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + + +def rewrite_logs(logs): + new_logs = {} + for k, v in logs.items(): + if isinstance(v, str): + continue + k = k.replace('/', '_') + if k.startswith('eval_'): + k = k[len('eval_'):] + k = f'eval/{k}' + elif k.startswith('test_'): + k = k[len('test_'):] + k = f'test/{k}' + else: + k = f'train/{k}' + new_logs[k] = v + return new_logs diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/wandb.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/wandb.py new file mode 100644 index 0000000000000000000000000000000000000000..9ea95ed39c466650d43f7667594b4b4c8d4c439a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/callbacks/wandb.py @@ -0,0 +1,31 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os + +from swift.utils import check_json_format, is_last_rank +from .base import MegatronCallback +from .utils import rewrite_logs + + +class WandbCallback(MegatronCallback): + + def __init__(self, trainer): + super().__init__(trainer) + args = self.args + self.config = check_json_format(vars(args)) + if args.wandb_exp_name is None: + args.wandb_exp_name = args.output_dir + self.save_dir = args.output_dir + self.writer = None + self.setup() + + def setup(self): + import wandb + args = self.args + if is_last_rank(): + wandb.init(dir=self.save_dir, name=args.wandb_exp_name, project=args.wandb_project, config=self.config) + self.writer = wandb + + def on_log(self, logs): + logs = rewrite_logs(logs) + if is_last_rank(): + self.writer.log(logs, step=self.state.iteration) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/convert.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/convert.py new file mode 100644 index 0000000000000000000000000000000000000000..b426d3414710fb66f93bc111f8b548efb5e24777 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/convert.py @@ -0,0 +1,120 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +import math +import os +import shutil +import torch +from transformers.utils import strtobool + +from swift.arguments import ExportArguments +from swift.pipelines import prepare_model_template +from swift.utils import get_logger, get_n_params_grads, is_master +from .arguments import MegatronArguments +from .model import get_mcore_model +from .utils import (load_mcore_checkpoint, patch_torch_dist_shard, prepare_mcore_model, save_mcore_checkpoint, + test_convert_precision) + +logger = get_logger() + +convert_kwargs = { + 'use_cpu_initialization': True, + 'no_save_optim': True, + 'no_save_rng': True, + 'no_load_optim': True, + 'no_load_rng': True, + 'finetune': True, + 'attention_backend': 'unfused', + 'padding_free': False, + 'recompute_granularity': 'none', # deepseek-v4 +} + + +def convert_hf2mcore(args: ExportArguments) -> None: + args.experts_impl = 'eager' # Compatible with transformers 5.4.0 + hf_model, template = prepare_model_template(args, patch_offload=not args.test_convert_precision) + processor = template.processor + if args.thread_count is None: + checkpoint_size = sum(get_n_params_grads(hf_model)[0]) * torch.finfo(args.torch_dtype).bits // 8e9 + args.thread_count = max(math.ceil(checkpoint_size / 10), 2) # 10GB + patch_torch_dist_shard(args.thread_count) + + hf_config = processor.model_info.config + current_convert_kwargs = convert_kwargs.copy() + if args.model_info.is_moe_model: + current_convert_kwargs['moe_grouped_gemm'] = True + megatron_args = MegatronArguments( + model=args.model, + model_type=args.model_type, + **current_convert_kwargs, + output_dir=args.output_dir, + torch_dtype=args.torch_dtype) + + mg_model = get_mcore_model(megatron_args, hf_config)[0] + logger.info('Megatron model created successfully.') + bridge = mg_model.config.bridge + bridge.load_weights([mg_model], args.model_info.model_dir) + logger.info('Successfully transferred HF model weights to MG model.') + _test_convert_precision = strtobool(os.getenv('SWIFT_TEST_CONVERT_PRECISION', '0')) + if not _test_convert_precision: + args.save_args() + logger.info('Saving the model...') + save_mcore_checkpoint(megatron_args, [mg_model]) + # Place it at the end to avoid test_convert_precision affecting precision. + if args.test_convert_precision: + test_convert_precision(megatron_args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype) + + +def convert_mcore2hf(args: ExportArguments) -> None: + args.experts_impl = 'eager' + _, template = prepare_model_template(args, load_model=False) + processor = template.processor + + hf_config = processor.model_info.config + current_convert_kwargs = convert_kwargs.copy() + if args.model_info.is_moe_model: + current_convert_kwargs['moe_grouped_gemm'] = True + extra_config = MegatronArguments.load_args_config(args.mcore_adapter or args.mcore_model) + extra_config['mcore_adapter'] = args.mcore_adapter + if args.mcore_model is not None: + extra_config['mcore_model'] = args.mcore_model + current_convert_kwargs.update(extra_config) + megatron_args = MegatronArguments( + model=args.model, + model_type=args.model_type, + **current_convert_kwargs, + output_dir=args.output_dir if args.to_mcore else None, + torch_dtype=args.torch_dtype) + + mg_model = get_mcore_model(megatron_args, hf_config)[0] + if megatron_args.mcore_model is None: + raise ValueError('Please specify `--mcore_model`.') + load_mcore_checkpoint(megatron_args, [mg_model], load_arg='mcore_model') + if megatron_args.mcore_adapter is not None: + peft_model = prepare_mcore_model(megatron_args, mg_model) + load_mcore_checkpoint(megatron_args, [mg_model], load_arg='mcore_adapter') + logger.info('Merge LoRA...') + mg_model = peft_model.merge_and_unload() + logger.info('Megatron model created successfully.') + if args.to_hf: + bridge = mg_model.config.bridge + logger.info('Converting weights and saving the model...') + bridge.save_weights([mg_model], args.output_dir, args=megatron_args, processor=processor) + if is_master(): + if args.ckpt_dir: + args_path = os.path.join(args.ckpt_dir, 'args.json') + shutil.copy(args_path, os.path.join(args.output_dir, 'args.json')) + else: + args.save_args(args.output_dir) + if args.test_convert_precision: + hf_model, template = prepare_model_template(args, model=args.output_dir) + test_convert_precision( + megatron_args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype) + elif args.to_mcore: + if args.thread_count is None: + checkpoint_size = sum(get_n_params_grads(mg_model)[0]) * torch.finfo(args.torch_dtype).bits // 8e9 + args.thread_count = max(math.ceil(checkpoint_size / 10), 2) # 10GB + patch_torch_dist_shard(args.thread_count) + + args.save_args() + logger.info('Saving the model...') + save_mcore_checkpoint(megatron_args, [mg_model]) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/init.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/init.py new file mode 100644 index 0000000000000000000000000000000000000000..973c1b01ad0c8213b01cf6c207d9b7725f8e1cdf --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/init.py @@ -0,0 +1,259 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import concurrent.futures +import inspect +import logging +import os +import torch +import torch.distributed as dist +from contextlib import contextmanager +from copy import copy, deepcopy +from tqdm import tqdm +from transformers.modeling_utils import custom_object_save +from transformers.utils import is_torch_npu_available +from typing import Union + +from swift.model import get_model_processor, save_checkpoint +from swift.utils import (HfConfigFactory, disable_safe_ddp_context_use_barrier, get_logger, get_modules_to_not_convert, + get_multimodal_target_regex, is_master, split_list) + +logger = get_logger() + + +def _patch__batched_p2p_ops(): + from megatron.core.pipeline_parallel import p2p_communication + + _batched_p2p_ops_origin = p2p_communication._batched_p2p_ops + + def _batched_p2p_ops(**kwargs): + kwargs['group'] = None + return _batched_p2p_ops_origin(**kwargs) + + p2p_communication._batched_p2p_ops = _batched_p2p_ops + + +def _patch_torch_FileSystemReader(): + from torch.distributed.checkpoint.filesystem import FileSystemReader + from torch.futures import Future + _origin_read_data = FileSystemReader.read_data + _origin__slice_file = FileSystemReader._slice_file + READER_MAX_WORKERS = int(os.environ.get('MCORE_READER_MAX_WORKERS', '16')) + + @contextmanager + def _patch__slice_file(prog_bar): + + def _slice_file(self, *args, **kwargs): + prog_bar.update() + return _origin__slice_file(self, *args, **kwargs) + + FileSystemReader._slice_file = _slice_file + try: + yield + finally: + FileSystemReader._slice_file = _origin__slice_file + + def read_data(self, plan, planner): + + def _worker(plan_shard): + _origin_read_data(self, plan_shard, planner) + + prog_bar = tqdm(total=len(plan.items), dynamic_ncols=True, desc='Loading: ') + plan_shards = split_list(plan.items, READER_MAX_WORKERS, contiguous=False) + with _patch__slice_file(prog_bar): + with concurrent.futures.ThreadPoolExecutor(max_workers=READER_MAX_WORKERS) as pool: + futures = [] + for i in range(READER_MAX_WORKERS): + plan_shard = copy(plan) + plan_shard.items = plan_shards[i] + futures.append(pool.submit(_worker, plan_shard)) + concurrent.futures.wait(futures) + prog_bar.close() + fut: Future = Future() + fut.set_result(None) + return fut + + FileSystemReader.read_data = read_data + + +def _dcp_validation_returns_errors(default_planner) -> bool: + """Whether `_validate_global_plan` is expected to return a list of error messages.""" + try: + # The caller is what defines the contract, so it is the most reliable thing to inspect. + source = inspect.getsource(default_planner.DefaultSavePlanner._create_global_plan) + return 'validation_errors' in source + except (OSError, TypeError): + pass + annotation = inspect.signature(default_planner._validate_global_plan).return_annotation + if annotation is inspect.Signature.empty: + logger.warning(f'Could not determine the `_validate_global_plan` contract of torch=={torch.__version__}; ' + 'assuming the legacy boolean form.') + return False + return annotation not in (bool, 'bool') + + +def _patch_validate_non_overlapping_shards_metadata(): + # too slow + from torch.distributed._shard.sharded_tensor import api + from torch.distributed._shard.sharding_spec import api as api2 + from torch.distributed.checkpoint import default_planner + + def validate_non_overlapping_shards_metadata(*args, **kwargs): + pass + + api.validate_non_overlapping_shards_metadata = validate_non_overlapping_shards_metadata + api2.validate_non_overlapping_shards_metadata = validate_non_overlapping_shards_metadata + + # The return contract changed across torch versions: it used to be a bool (falsy meaning + # "invalid"), while newer versions return a list of error messages (empty meaning "valid"). + # Returning the wrong type is not harmless -- a bool sends the newer caller into its error + # branch, where `'; '.join(True)` raises `TypeError: can only join an iterable` and buries the + # real reason for the failure. + if _dcp_validation_returns_errors(default_planner): + + def _validate_global_plan(*args, **kwargs): + return [] + else: + + def _validate_global_plan(*args, **kwargs): + return True + + default_planner._validate_global_plan = _validate_global_plan + + +def _patch_unified_memory(): + if is_torch_npu_available(): + return + + from torch.utils import cpp_extension + load_inline = cpp_extension.load_inline + + def _new_load_inline(*args, **kwargs): + name = kwargs.get('name') + if name == 'managed_alloc_runtime': + raise RuntimeError + return load_inline(*args, **kwargs) + + # not create unified memory mempool + cpp_extension.load_inline = _new_load_inline + try: + from megatron.core.inference import unified_memory + except Exception: + pass + finally: + cpp_extension.load_inline = load_inline + + +def _patch_mcore_bridge(): + import mcore_bridge + from mcore_bridge import GPTBridge + logger.info(f'mcore_bridge.__version__: {mcore_bridge.__version__}') + origin_save_weights = GPTBridge.save_weights + + def save_weights( + self, + mg_models, + output_dir: str, + peft_format: bool = False, + max_shard_size: str = '5GB', + args=None, + processor=None, + save_missing_weights: Union[bool, str] = False, + ) -> None: + origin_save_weights( + self, + mg_models, + output_dir, + peft_format=peft_format, + max_shard_size=max_shard_size, + save_missing_weights=save_missing_weights) + if processor is None or args is None: + return + hf_config = self.config.hf_config + hf_config = deepcopy(hf_config) + if is_master() and not hasattr(self, 'hf_model'): + if hasattr(self, 'get_hf_meta_model'): + self.hf_model = self.get_hf_meta_model() + self.hf_model.model_meta = processor.model_meta + self.hf_model.model_info = processor.model_info + else: + with torch.device('meta'), disable_safe_ddp_context_use_barrier(): + self.hf_model = get_model_processor( + args.model_dir, model_type=args.model_type, return_dummy_model=True)[0] + + if is_master(): + if peft_format: + peft_config = copy(mg_models[0].peft_config[self._adapter_name]) + if self.config.task_type == 'seq_cls': + peft_config.task_type = 'SEQ_CLS' + if self.is_multimodal and 'all-linear' in args.target_modules: + peft_config.target_modules = get_multimodal_target_regex( + self.hf_model, + freeze_llm=args.freeze_llm, + freeze_vit=args.freeze_vit, + freeze_aligner=args.freeze_aligner, + include_embedding='all-embedding' in args.target_modules, + exclude_router='all-router' not in args.target_modules) + else: + assert not isinstance(peft_config.target_modules, str), ( + 'target_regex is not currently supported for LoRA conversion. Please set `--merge_lora true`.') + peft_config.target_modules = self._peft_target_modules + peft_config.modules_to_save = self._peft_modules_to_save + peft_config.save_pretrained(output_dir) + else: + config = self.config + llm_config = HfConfigFactory.get_text_config(hf_config) + if config.mtp_num_layers: + for key in ['num_nextn_predict_layers', 'mtp_num_hidden_layers']: + if hasattr(llm_config, key): + setattr(llm_config, key, config.mtp_num_layers) + break + else: + llm_config.num_nextn_predict_layers = config.mtp_num_layers + HfConfigFactory.del_config_attr(hf_config, 'quantization_config') + expert_dtype = None + if config.fp8 is not None and config.fp8_recipe == 'blockwise' and config.fp8_param: + from transformers.utils.quantization_config import FineGrainedFP8Config + modules_to_not_convert = get_modules_to_not_convert(self.hf_model) + if hasattr(self, '_fp8_skip_modules'): + modules_to_not_convert = (modules_to_not_convert or []) + list(self._fp8_skip_modules) + hf_config.quantization_config = FineGrainedFP8Config(modules_to_not_convert=modules_to_not_convert) + expert_dtype = 'fp8' + if args.model_type == 'deepseek_v4': + HfConfigFactory.set_config_attr(hf_config, 'expert_dtype', expert_dtype) + hf_config.save_pretrained(output_dir) + if getattr(self.hf_model, '_auto_class') is not None: + try: + custom_object_save(self.hf_model, output_dir, config=hf_config) + except FileNotFoundError as e: + logger.error(f'custom_object_save Error: {e}') + save_checkpoint( + None, + processor, + output_dir, + model_dirs=[args.model_dir], + additional_saved_files=self.hf_model.model_meta.additional_saved_files) + logger.info(f'Successfully saved `safetensors` model weights in `{output_dir}`.') + dist.barrier() # Ensure all weights are saved completely + + GPTBridge.save_weights = save_weights + + +def init_megatron_env(): + os.environ.pop('VLLM_USE_MODELSCOPE', None) + logging_level = logging.root.level + _patch_unified_memory() + if is_torch_npu_available(): + from swift.model.npu_patcher import patch_mindspeed_fla_gdn_implementation + patch_mindspeed_fla_gdn_implementation() + _patch__batched_p2p_ops() + logging.root.setLevel(logging_level) # revert logger level + try: + _patch_torch_FileSystemReader() + except Exception: + logger.warning('Failed to patch FileSystemReader.') + try: + _patch_validate_non_overlapping_shards_metadata() + except Exception: + logger.warning('Patch validate_non_overlapping_shards_metadata failed.') + pass + import megatron.core + logger.info(f'megatron.core.__version__: {megatron.core.__version__}') diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/model/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/model/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eb5d16ad4fc7ea1571ded4e70189a5a4531b2911 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/model/__init__.py @@ -0,0 +1 @@ +from .utils import get_mcore_model diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/model/utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/model/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1d01ab933ad918dfd7500b04f225466bc44a4edd --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/model/utils.py @@ -0,0 +1,293 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from dataclasses import fields +from mcore_bridge import ModelConfig +from mcore_bridge import get_mcore_model as _get_mcore_model +from mcore_bridge import hf_to_mcore_config +from transformers.utils import is_torch_npu_available +from typing import Any, Generator, Optional, Tuple + +from swift.utils import get_logger + +logger = get_logger() + + +def _check_attention_backend(args, config): + """Validate attention backend compatibility with configuration.""" + attention_backend = config.attention_backend.name + if attention_backend == 'flash' and config.softmax_type == 'learnable': + raise ValueError(f'Attention backend "{attention_backend}" does not support learnable softmax_type.') + + +def _check_padding_free(args, config): + """Validate and adjust padding_free setting based on configuration constraints.""" + if not args.padding_free: + return + + attention_backend = config.attention_backend.name + message = None + + if attention_backend == 'unfused': + message = f'Attention backend "{attention_backend}" is not supported in padding-free mode' + + if message: + logger.warning(f'{message}. Setting args.padding_free to False.') + args.padding_free = False + + +def get_mcore_model_config(args, hf_config): + kwargs = hf_to_mcore_config(hf_config) + kwargs['mcore_model_type'] = args.megatron_model_meta.model_type + kwargs['hf_config'] = hf_config + for f in fields(ModelConfig): + key, value = f.name, getattr(args, f.name, None) + if value is None or isinstance(value, (list, tuple)) and len(value) == 0: + continue + kwargs[key] = value + + if args.task_type == 'seq_cls': + args.problem_type = args.problem_type or getattr(hf_config, 'problem_type', None) + logger.info(f'args.problem_type: {args.problem_type}') + + kwargs['params_dtype'] = args.torch_dtype + kwargs['num_layers_in_first_pipeline_stage'] = args.decoder_first_pipeline_num_layers + kwargs['num_layers_in_last_pipeline_stage'] = args.decoder_last_pipeline_num_layers + kwargs['fp4_param'] = args.fp4_param_gather + kwargs['fp8_param'] = args.fp8_param_gather + swiglu = kwargs.get('swiglu', True) + add_bias_linear = kwargs.get('add_bias_linear', False) + num_moe_experts = kwargs.get('num_moe_experts', None) + position_embedding_type = kwargs.get('position_embedding_type', 'rope') + if position_embedding_type != 'rope': + kwargs['apply_rope_fusion'] = False + if not swiglu and not add_bias_linear: + kwargs['bias_activation_fusion'] = False + if add_bias_linear and num_moe_experts and args.moe_grouped_gemm: + kwargs['bias_dropout_fusion'] = False + if num_moe_experts is None: + kwargs['expert_model_parallel_size'] = 1 + kwargs['expert_tensor_parallel_size'] = 1 + + if args.router_replay_mode != 'disabled': + kwargs['moe_enable_routing_replay'] = True + if args.megatron_extra_kwargs: + kwargs.update(args.megatron_extra_kwargs) + config = ModelConfig(**kwargs) + if is_torch_npu_available() and getattr(args, 'attention_backend', 'flash') != 'local': + setattr(config, 'use_flash_attn', True) + _check_attention_backend(args, config) + _check_padding_free(args, config) + return config + + +class MegatronBridgeBackend: + """Adapter for NVIDIA ``megatron.bridge.AutoBridge``. + + Limitations: + - LoRA / PEFT loading is not yet supported. + - MLLM is not yet supported + - FP8 export is not yet supported. + """ + + def __init__(self, auto_bridge: Any, hf_config: Optional[Any] = None): + self._bridge = auto_bridge + self._hf_config = hf_config + + @classmethod + def from_hf_config(cls, hf_config) -> 'MegatronBridgeBackend': + from megatron.bridge.models.conversion.auto_bridge import AutoBridge + return cls(AutoBridge.from_hf_config(hf_config), hf_config) + + def load_weights(self, models, hf_model_dir, peft_format=False, adapter_name='default', converter=None): + if peft_format: + raise NotImplementedError('LoRA loading via megatron-bridge backend is not yet supported. ' + 'Please use bridge_backend="mcore-bridge" for LoRA training.') + if converter is not None: + logger.warning('converter is not supported by megatron-bridge backend, ignoring.') + self._bridge.load_hf_weights(models, hf_path=hf_model_dir) + + def export_weights(self, + models, + target_device=None, + only_master_rank=False, + peft_format=False, + adapter_name='default', + converter=None, + tqdm_desc='Exporting: ', + disable_tqdm=True, + _is_saving=False) -> Generator[Tuple[str, 'torch.Tensor'], None, None]: + if peft_format: + raise NotImplementedError('LoRA export via megatron-bridge backend is not yet supported. ' + 'Please use bridge_backend="mcore-bridge" for LoRA training.') + cpu = (target_device == 'cpu') + for weight_tuple in self._bridge.export_hf_weights(models, cpu=cpu): + key = weight_tuple.param_name + tensor = weight_tuple.weight + if converter is not None and tensor is not None: + kv = converter(key, tensor) + if kv is None: + continue + key, tensor = kv + yield key, tensor + + def save_weights(self, + models, + output_dir, + peft_format=False, + max_shard_size='5GB', + args=None, + processor=None) -> None: + if peft_format: + raise NotImplementedError('LoRA saving via megatron-bridge backend is not yet supported. ' + 'Please use bridge_backend="mcore-bridge" for LoRA training.') + + # 1. Save weights via megatron-bridge (safetensors format) + self._bridge.save_hf_weights(models, path=output_dir) + + # 2. Save HF config and tokenizer on rank 0. + # We use the original HF config (not the one reconstructed by megatron-bridge) + # because the bridge's config-only path may drop fields like num_attention_heads. + import torch.distributed as dist + is_master = (not dist.is_initialized()) or dist.get_rank() == 0 + if is_master and args is not None and self._hf_config is not None: + from copy import deepcopy + + from swift.model import save_checkpoint + from swift.utils import HfConfigFactory + hf_config = deepcopy(self._hf_config) + llm_config = HfConfigFactory.get_text_config(hf_config) + + # MTP: write back num_nextn_predict_layers + mtp_num_layers = getattr(args, 'mtp_num_layers', None) + if mtp_num_layers: + for key in ['num_nextn_predict_layers', 'mtp_num_hidden_layers']: + if hasattr(llm_config, key): + setattr(llm_config, key, mtp_num_layers) + break + else: + llm_config.num_nextn_predict_layers = mtp_num_layers + + HfConfigFactory.del_config_attr(hf_config, 'quantization_config') + + # FP8: write back quantization_config + expert_dtype = None + fp8_format = getattr(args, 'fp8_format', None) + fp8_recipe = getattr(args, 'fp8_recipe', 'delayed') + fp8_param = getattr(args, 'fp8_param_gather', False) + if fp8_format is not None and fp8_recipe == 'blockwise' and fp8_param: + from transformers.utils.quantization_config import FineGrainedFP8Config + hf_config.quantization_config = FineGrainedFP8Config() + expert_dtype = 'fp8' + if getattr(args, 'model_type', None) == 'deepseek_v4': + HfConfigFactory.set_config_attr(hf_config, 'expert_dtype', expert_dtype) + + hf_config.save_pretrained(output_dir) + if processor is not None: + additional_saved_files = getattr(getattr(processor, 'model_meta', None), 'additional_saved_files', None) + save_checkpoint( + None, + processor, + output_dir, + model_dirs=[args.model_dir], + additional_saved_files=additional_saved_files) + if dist.is_initialized(): + dist.barrier() + + +def get_mcore_model(args, hf_config): + bridge_backend = args.bridge_backend + if bridge_backend == 'megatron-bridge': + return _get_megatron_bridge_model(args, hf_config) + config = get_mcore_model_config(args, hf_config) + models = _get_mcore_model(config) + + return models + + +def _get_megatron_bridge_model(args, hf_config): + import dataclasses + + backend = MegatronBridgeBackend.from_hf_config(hf_config) + auto_bridge = backend._bridge + + # Validate model support via AutoBridge.supports() + from megatron.bridge.models.conversion.auto_bridge import AutoBridge + if not AutoBridge.supports(hf_config): + raise ValueError(f'Model {getattr(hf_config, "model_type", "unknown")} is not supported by ' + f'megatron-bridge. Please use bridge_backend="mcore-bridge" or check ' + f'AutoBridge.list_supported_models() for supported architectures.') + + # --- Step 1: Get provider (GPTModelProvider, which extends TransformerConfig) --- + provider = auto_bridge.to_megatron_provider(load_weights=False) + + # --- Step 2: Build overrides from args --- + # Auto-match: iterate over provider's dataclass fields and pick up matching args fields. + # This mirrors mcore-bridge's get_mcore_model_config which does: + # for f in fields(ModelConfig): kwargs[f.name] = getattr(args, f.name, None) + overrides = {} + provider_fields = {f.name for f in dataclasses.fields(provider)} + for field_name in provider_fields: + value = getattr(args, field_name, None) + if value is None or (isinstance(value, (list, tuple)) and len(value) == 0): + continue + overrides[field_name] = value + + # Explicit field name mappings (args name → provider field name) + explicit_mappings = { + 'decoder_first_pipeline_num_layers': 'num_layers_in_first_pipeline_stage', + 'decoder_last_pipeline_num_layers': 'num_layers_in_last_pipeline_stage', + } + for args_key, provider_key in explicit_mappings.items(): + value = getattr(args, args_key, None) + if value is not None and provider_key in provider_fields: + overrides[provider_key] = value + + # dtype + dtype = getattr(args, 'torch_dtype', None) + + # MoE: if no experts, force EP/ETP to 1 + # num_moe_experts comes from HF config (parsed by AutoBridge into provider), + # not from args — so check provider too. + num_moe_experts = overrides.get('num_moe_experts') or getattr(provider, 'num_moe_experts', None) + if num_moe_experts is None: + overrides['expert_model_parallel_size'] = 1 + overrides['expert_tensor_parallel_size'] = 1 + + # Router replay + if getattr(args, 'router_replay_mode', 'disabled') != 'disabled': + if 'moe_enable_routing_replay' in provider_fields: + overrides['moe_enable_routing_replay'] = True + + # megatron_extra_kwargs (user-specified raw overrides) + if getattr(args, 'megatron_extra_kwargs', None): + overrides.update(args.megatron_extra_kwargs) + + # padding_free requires variable_seq_lengths=True so that RotaryEmbedding + # generates freqs matching the actual packed sequence length (cu_seqlens[-1]) + # instead of the fixed seq_length. Without this, mcore-bridge's patcher + # use_batched_rope check fails and falls back to the original + # _apply_rotary_pos_emb_thd which calls torch.split on a padded tensor. + if getattr(args, 'padding_free', False) and 'variable_seq_lengths' in provider_fields: + overrides['variable_seq_lengths'] = True + + # --- Step 3: Apply overrides and finalize --- + provider.apply_overrides_and_finalize(dtype=dtype, overrides=overrides) + + import torch.nn.functional as F + provider.swiglu = (provider.gated_linear_unit and provider.activation_func is F.silu) + + # --- Step 4: Create raw models (no DDP/Float16 wrapping) --- + # swift's wrap_model handles DDP/Float16 wrapping with the correct DDP config from args. + models = provider.provide_distributed_model( + wrap_with_ddp=False, + mixed_precision_wrapper=None, + use_cpu_initialization=getattr(args, 'use_cpu_initialization', False), + ) + if not isinstance(models, list): + models = [models] + + # --- Step 5: Attach backend to model.config.bridge --- + for model in models: + model.config.bridge = backend + + logger.info('Created Megatron model via megatron-bridge backend') + return models diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2a2c44e567d8f99662512eb039e769a2eb07a520 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING + +from swift.utils.import_utils import _LazyModule + +if TYPE_CHECKING: + from .export import megatron_export_main + from .train import megatron_pretrain_main, megatron_rlhf_main, megatron_sft_main +else: + _import_structure = { + 'train': ['megatron_pretrain_main', 'megatron_rlhf_main', 'megatron_sft_main'], + 'export': ['megatron_export_main'], + } + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/export/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/export/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..97590932f0ebfbde41a5c9f145a24b10ba6e658b --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/export/__init__.py @@ -0,0 +1 @@ +from .export import megatron_export_main diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/export/export.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/export/export.py new file mode 100644 index 0000000000000000000000000000000000000000..ac616448a140e8ca66dbf4ad7c29bdb02511e930 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/export/export.py @@ -0,0 +1,135 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import shutil +import torch.distributed as dist +from transformers.utils import strtobool +from typing import List, Optional, Union + +from swift.megatron.arguments import MegatronExportArguments +from swift.megatron.convert import test_convert_precision +from swift.megatron.model import get_mcore_model +from swift.megatron.utils import load_mcore_checkpoint, prepare_mcore_model, save_mcore_checkpoint +from swift.pipelines import SwiftPipeline, prepare_model_template +from swift.utils import disable_safe_ddp_context_use_barrier, get_logger, is_master + +logger = get_logger() + + +class MegatronExport(SwiftPipeline): + args_class = MegatronExportArguments + args: args_class + + def run(self): + os.environ['DISABLE_MP_DDP'] = 'true' + args = self.args + if args.to_hf: + self.convert_mcore2hf() + elif args.to_mcore: + self.convert_hf2mcore() + + def convert_mcore2hf(self) -> None: + args = self.args + args.experts_impl = 'eager' + download_model = args.model is not None + _, template = prepare_model_template(args, load_model=False, download_model=download_model) + self.processor = template.processor + hf_config = self.processor.model_info.config + mg_model = get_mcore_model(args, hf_config)[0] + logger.info('Megatron model created successfully.') + bridge = mg_model.config.bridge + if args.mcore_model is not None: + load_mcore_checkpoint(args, [mg_model], load_arg='mcore_model') + elif args.model is not None: + bridge.load_weights([mg_model], args.model_info.model_dir) + else: + raise ValueError('Please specify `--mcore_model` or `--model`.') + if args.adapters or args.mcore_adapter is not None: + peft_model = prepare_mcore_model(args, mg_model) + if args.mcore_adapter is not None: + load_mcore_checkpoint(args, [mg_model], load_arg='mcore_adapter') + elif args.adapters: + assert len(args.adapters) == 1, 'Currently only support one adapter' + bridge.load_weights([mg_model], args.adapters[0], peft_format=True) + if args.merge_lora: + logger.info('Merge LoRA...') + mg_model = peft_model.merge_and_unload() + logger.info('Converting weights and saving the model...') + save_peft_format = args.tuner_type == 'lora' and not args.merge_lora + # `to_hf` may load from a mcore checkpoint, so pass the source HF dir explicitly. + save_missing_weights = (args.save_missing_weights and args.model_info is not None and args.model_info.model_dir) + bridge.save_weights( + [mg_model], + args.output_dir, + peft_format=save_peft_format, + args=args, + processor=self.processor, + save_missing_weights=save_missing_weights, + ) + if is_master(): + if args.ckpt_dir: + args_path = os.path.join(args.ckpt_dir, 'args.json') + shutil.copy(args_path, os.path.join(args.output_dir, 'args.json')) + else: + args.save_args(args.output_dir) + if args.test_convert_precision: + with disable_safe_ddp_context_use_barrier(): + if save_peft_format: + kwargs = {'adapters': [args.output_dir]} + else: + kwargs = {'model': args.output_dir, 'torch_dtype': None, 'adapters': []} + device_map = args.device_map or 'auto' + hf_model, template = prepare_model_template( + args, device_map=device_map, **kwargs) if is_master() else (None, template) + test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype) + dist.barrier() + + def convert_hf2mcore(self) -> None: + args = self.args + args.experts_impl = 'eager' + download_model = args.model is not None + _, template = prepare_model_template(args, load_model=False, download_model=download_model) + self.processor = template.processor + hf_config = self.processor.model_info.config + mg_model = get_mcore_model(args, hf_config)[0] + logger.info('Megatron model created successfully.') + bridge = mg_model.config.bridge + if args.model is not None: + bridge.load_weights([mg_model], args.model_info.model_dir) + elif args.mcore_model is not None: + load_mcore_checkpoint(args, [mg_model], load_arg='mcore_model') + else: + raise ValueError('Please specify `--mcore_model` or `--model`.') + dist.barrier() + if args.adapters or args.mcore_adapter is not None: + peft_model = prepare_mcore_model(args, mg_model) + if args.adapters: + assert len(args.adapters) == 1, 'Currently only support one adapter' + bridge.load_weights([mg_model], args.adapters[0], peft_format=True) + elif args.mcore_adapter is not None: + load_mcore_checkpoint(args, [mg_model], load_arg='mcore_adapter') + if args.merge_lora: + logger.info('Merge LoRA...') + mg_model = peft_model.merge_and_unload() + logger.info('Successfully transferred HF model weights to MG model.') + _test_convert_precision = strtobool(os.getenv('SWIFT_TEST_CONVERT_PRECISION', '0')) + if not _test_convert_precision: + args.save_args(args.output_dir) + logger.info('Saving the model...') + save_peft_format = args.tuner_type == 'lora' and not args.merge_lora + save_mcore_checkpoint(args, [mg_model], peft_format=save_peft_format) + # hf_model does not support loading args.mcore_adapter, so test_convert_precision cannot be performed + support_convert_precision = args.mcore_adapter is None + if args.test_convert_precision: + if support_convert_precision: + with disable_safe_ddp_context_use_barrier(): + device_map = args.device_map or 'auto' + hf_model, template = prepare_model_template( + args, device_map=device_map) if is_master() else (None, template) + test_convert_precision(args, hf_model, mg_model, template, test_convert_dtype=args.test_convert_dtype) + dist.barrier() + else: + logger.warning('Skip test_convert_precision because `--mcore_adapter` is specified.') + + +def megatron_export_main(args: Optional[Union[List[str], MegatronExportArguments]] = None): + return MegatronExport(args).main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5cb351d2484c7e484a2711d50136e03c515c46e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/__init__.py @@ -0,0 +1,24 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING + +from swift.utils.import_utils import _LazyModule + +if TYPE_CHECKING: + from .pretrain import megatron_pretrain_main + from .rlhf import megatron_rlhf_main + from .sft import megatron_sft_main +else: + _import_structure = { + 'pretrain': ['megatron_pretrain_main'], + 'rlhf': ['megatron_rlhf_main'], + 'sft': ['megatron_sft_main'], + } + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/pretrain.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/pretrain.py new file mode 100644 index 0000000000000000000000000000000000000000..c62b9b806b12744cbbaf20b8f07f2725de1baf11 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/pretrain.py @@ -0,0 +1,17 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import List, Optional, Union + +from swift.megatron.arguments import MegatronPretrainArguments +from swift.utils import get_logger +from .sft import MegatronSft + +logger = get_logger() + + +class MegatronPretrain(MegatronSft): + args_class = MegatronPretrainArguments + args: args_class + + +def megatron_pretrain_main(args: Optional[Union[List[str], MegatronPretrainArguments]] = None): + return MegatronPretrain(args).main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/rlhf.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/rlhf.py new file mode 100644 index 0000000000000000000000000000000000000000..05e3aef57064f5ae3d3b85c344dc75b6b0b1d393 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/rlhf.py @@ -0,0 +1,73 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import importlib +from typing import List, Optional, Union + +from swift.megatron.arguments import MegatronRLHFArguments +from swift.pipelines.train import prepare_kto_dataset +from swift.utils import get_current_device, get_logger, is_last_rank +from .sft import MegatronSft + +logger = get_logger() + + +class MegatronRLHF(MegatronSft): + args_class = MegatronRLHFArguments + args: args_class + + def prepare_trainer(self): + args = self.args + trainer_mapping = { + 'dpo': 'MegatronDPOTrainer', + 'gkd': 'MegatronGKDTrainer', + 'grpo': 'MegatronGRPOTrainer', + 'kto': 'MegatronKTOTrainer', + 'rm': 'MegatronRewardTrainer' + } + module = importlib.import_module('swift.megatron.trainers') + trainer_cls_name = trainer_mapping.get(args.rlhf_type) + if trainer_cls_name is None: + raise ValueError(f'The current Megatron-SWIFT does not support rlhf_type: {args.rlhf_type}.') + trainer_cls = getattr(module, trainer_cls_name) + kwargs = {} + if args.rlhf_type in ('grpo', 'gkd'): + kwargs['vllm_client'] = self._prepare_vllm_client() + return trainer_cls(args, self.template, **kwargs) + + def _prepare_template(self) -> None: + super()._prepare_template() + model_mapping = {'grpo': 'train', 'gkd': 'train', 'kto': 'kto'} + self.template.set_mode(model_mapping.get(self.args.rlhf_type, 'rlhf')) + + def _get_dataset(self): + args = self.args + train_dataset, val_dataset = super()._get_dataset() + if args.rlhf_type == 'kto': + train_dataset, val_dataset = prepare_kto_dataset(args, train_dataset, val_dataset) + return train_dataset, val_dataset + + def _prepare_vllm_client(self): + # Only prepare vLLM client for server mode + if self.args.rlhf_type not in ('grpo', 'gkd') or self.args.vllm_mode != 'server': + return None + # GKD may not use vLLM (off-policy mode) + if not getattr(self.args, 'use_vllm', False): + return None + + from swift.rlhf_trainers.vllm_client import VLLMClient + vllm_client = None + if is_last_rank(): + logger.info('Start connecting to vLLM server') + vllm_client = VLLMClient( + base_urls=self.args.vllm_server_base_url, + hosts=self.args.vllm_server_host, + server_ports=self.args.vllm_server_port, + group_ports=self.args.vllm_server_group_port, + connection_timeout=self.args.vllm_server_timeout) + vllm_client.close_communicator() + vllm_client.init_communicator(device=get_current_device()) + logger.info('Connected to vLLM server') + return vllm_client + + +def megatron_rlhf_main(args: Optional[Union[List[str], MegatronRLHFArguments]] = None): + return MegatronRLHF(args).main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/sft.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/sft.py new file mode 100644 index 0000000000000000000000000000000000000000..e3fe6a1ffb4acb1f95b94839a4dc76a090d98110 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/pipelines/train/sft.py @@ -0,0 +1,93 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import torch +import torch.distributed as dist +from dataclasses import asdict +from transformers.utils import is_torch_npu_available +from typing import List, Optional, Union + +from swift.megatron.arguments import MegatronSftArguments +from swift.megatron.trainers import MegatronEmbeddingTrainer, MegatronRerankerTrainer, MegatronTrainer +from swift.pipelines import SwiftSft +from swift.utils import append_to_jsonl, get_logger, is_last_rank, plot_images + +if is_torch_npu_available(): + # Enable Megatron on Ascend NPU + from swift.model.npu_patcher import apply_mindspeed_patches +else: + apply_mindspeed_patches = None + +logger = get_logger() + + +class MegatronSft(SwiftSft): + args_class = MegatronSftArguments + args: args_class + + def prepare_trainer(self): + args = self.args + if args.task_type == 'embedding': + return MegatronEmbeddingTrainer(self.args, self.template) + elif args.task_type in {'reranker', 'generative_reranker'}: + return MegatronRerankerTrainer(self.args, self.template) + else: + return MegatronTrainer(self.args, self.template) + + def _set_seed(self): + pass + + def __init__(self, args: Optional[Union[List[str], MegatronSftArguments]] = None) -> None: + self.train_msg = {} + super(SwiftSft, self).__init__(args) + args = self.args + if apply_mindspeed_patches is not None: + megatron_args = asdict(self.args) + if args.attention_backend != 'local': + # MindSpeed requires passing `use_flash_attn` to Megatron + # to enable flash attention on Ascend NPU. + args.use_flash_attn = True + megatron_args['use_flash_attn'] = True + apply_mindspeed_patches(megatron_args) + template_cls = args.template_meta.template_cls + if args.model_meta.is_multimodal and template_cls and template_cls.use_model: + kwargs = {'return_dummy_model': True} + else: + kwargs = {'load_model': False} + with torch.device('meta'): + self.model, self.processor = args.get_model_processor(**kwargs, download_model=args.mcore_model is None) + self._prepare_template() + args.save_args(args.output_dir) + self.template.use_megatron = True + + def run(self): + args = self.args + train_dataset, val_dataset = self._prepare_dataset() + args.init_iters(train_dataset, val_dataset) + trainer = self.prepare_trainer() + try: + trainer.train(train_dataset, val_dataset) + finally: + state = trainer.state + self._handle_trainer_state(trainer, is_last_rank()) + self.train_msg.update({ + 'last_model_checkpoint': state.last_model_checkpoint, + 'best_model_checkpoint': state.best_model_checkpoint, + 'best_metric': state.best_metric, + }) + # Visualization + if is_last_rank(): + images_dir = os.path.join(args.output_dir, 'images') + logger.info(f'images_dir: {images_dir}') + plot_images(images_dir, args.tensorboard_dir) + + jsonl_path = os.path.join(args.output_dir, 'logging.jsonl') + append_to_jsonl(jsonl_path, self.train_msg, strict=False, write_on_rank='last') + # Exceptions may cause the process to hang, preventing the exception from being propagated. + # Therefore, destroy_process_group() should not be placed inside the finally block. + if dist.is_initialized(): + dist.destroy_process_group() + return self.train_msg + + +def megatron_sft_main(args: Optional[Union[List[str], MegatronSftArguments]] = None): + return MegatronSft(args).main() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f49bffd0fee6412af90fd236447e91cb7f007e8d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/__init__.py @@ -0,0 +1,38 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from typing import TYPE_CHECKING + +from swift.utils.import_utils import _LazyModule + +if TYPE_CHECKING: + from .base import BaseMegatronTrainer + from .dpo_trainer import MegatronDPOTrainer + from .embedding_trainer import MegatronEmbeddingTrainer + from .gkd_trainer import MegatronGKDTrainer + from .grpo_trainer import MegatronGRPOTrainer + from .kto_trainer import MegatronKTOTrainer + from .reranker_trainer import MegatronRerankerTrainer + from .reward_trainer import MegatronRewardTrainer + from .rollout_mixin import MegatronRolloutMixin + from .trainer import MegatronTrainer +else: + _import_structure = { + 'dpo_trainer': ['MegatronDPOTrainer'], + 'gkd_trainer': ['MegatronGKDTrainer'], + 'grpo_trainer': ['MegatronGRPOTrainer'], + 'kto_trainer': ['MegatronKTOTrainer'], + 'reward_trainer': ['MegatronRewardTrainer'], + 'rollout_mixin': ['MegatronRolloutMixin'], + 'embedding_trainer': ['MegatronEmbeddingTrainer'], + 'reranker_trainer': ['MegatronRerankerTrainer'], + 'trainer': ['MegatronTrainer'], + 'base': ['BaseMegatronTrainer'], + } + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/base.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..048bf4dcf58539b07a6ed60474544b77299ecc2c --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/base.py @@ -0,0 +1,1066 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import dataclasses +import logging +import megatron.core +import operator +import os +import shutil +import torch +import torch.nn +from abc import ABC, abstractmethod +from contextlib import contextmanager, nullcontext +from functools import partial +from mcore_bridge import LoraParallelLinear +from megatron.core import mpu +from megatron.core.distributed import DistributedDataParallel as DDP +from megatron.core.distributed import FullyShardedDataParallel as megatron_FSDP +from megatron.core.distributed import finalize_model_grads +from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer +from megatron.core.pipeline_parallel import get_forward_backward_func +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.moe.moe_utils import track_moe_metrics +from megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper +from modelscope import check_local_model_is_latest +from packaging import version +from pathlib import Path +from typing import Callable, Dict, List, Optional + +from swift.dataset import RowPreprocessor +from swift.megatron.callbacks import megatron_callbacks_map +from swift.megatron.model import get_mcore_model +from swift.megatron.utils import (apply_router_replay_patch, disable_forward_pre_hook, enable_forward_pre_hook, + get_optimizer_param_scheduler, get_padding_to, init_persistent_async_worker, + initialize_tp_communicators, load_mcore_checkpoint, + logical_and_across_model_parallel_group, maybe_finalize_async_save, + prepare_mcore_model, reconstruct_tensor_cp, + reduce_max_stat_across_model_parallel_group, save_mcore_checkpoint, + should_disable_forward_pre_hook, warmup_jit_function, wrap_model) +from swift.template import Template +from swift.trainers import dynamic_gradient_checkpointing +from swift.trainers.utils import patch_modelscope_hub_timeout +from swift.utils import (deep_getattr, gc_collect, get_current_device, get_last_valid_indices, get_logger, is_last_rank, + is_master, ms_logger_context) +from .batch_sampler import MegatronPretrainingRandomSampler, MegatronPretrainingSampler +from .utils import TrainerState, build_streaming_dataloader, prepare_batch + +try: + from megatron.core.optimizer import param_group_identifier_keys + from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction +except ImportError: + param_group_identifier_keys = None + RouterReplay = None + RouterReplayAction = None + +mcore_016 = version.parse(megatron.core.__version__) >= version.parse('0.16.0rc0') + +logger = get_logger() + + +class BaseMegatronTrainer(ABC): + + def __init__(self, args, template: Template): + # validate mcore version and patch routing_replay + self.enable_routing_replay = args.router_replay_mode != 'disabled' + if self.enable_routing_replay: + apply_router_replay_patch() + + self.args = args + self.template = template + self.prepare_model() + # Sync template.padding_free after prepare_model(), because _check_padding_free + # may override args.padding_free for certain models (e.g. DSA attention). + if template.padding_free != args.padding_free: + logger.warning(f'template.padding_free({template.padding_free}) != args.padding_free({args.padding_free}), ' + f'syncing template.padding_free to {args.padding_free}.') + template.padding_free = args.padding_free + self.optimizer, self.opt_param_scheduler = self.get_optimizer_and_scheduler() + self.data_collator = self._get_data_collator() + + self.state = TrainerState(max_steps=args.train_iters) + initialize_embedding = args.new_special_tokens or args.task_type == 'seq_cls' + if initialize_embedding: + for m in self.unwrapped_models: + self._initialize_embedding(m) + self._load_checkpoint() + + self.eval_metrics = None + if args.check_model and hasattr(args, 'model_dir'): + with ms_logger_context(logging.CRITICAL), patch_modelscope_hub_timeout(): + config_info = self._collect_config_info() + config_info.update({ + 'invoked_by': 'local_trainer', + 'third_party': 'swift', + 'trainer_class': self.__class__.__name__, + 'trainer_backend': 'megatron', + }) + check_local_model_is_latest(args.model_info.model_dir, user_agent=config_info) + + self.callbacks = [] + for callback in args.callbacks: + self.callbacks.append(megatron_callbacks_map[callback](self)) + + if args.tp_comm_overlap: + initialize_tp_communicators(args, self.config) + + warmup_jit_function(self.config, args) + + if args.async_save and args.use_persistent_ckpt_worker: + init_persistent_async_worker() + + def _load_checkpoint(self): + args = self.args + if not args.finetune: + self.state.iteration = self._load_iteration() + if args.mcore_model is not None: + self.state.iteration = load_mcore_checkpoint( + args, self.wrapped_models, self.optimizer, self.opt_param_scheduler, load_arg='mcore_model') + if args.mcore_adapter is not None: + self.state.iteration = load_mcore_checkpoint( + args, self.wrapped_models, self.optimizer, self.opt_param_scheduler, load_arg='mcore_adapter') + self.state.consumed_train_samples = getattr(args, 'consumed_train_samples', 0) + + def call_event(self, event, **kwargs): + for callback in self.callbacks: + getattr(callback, event)(**kwargs) + + def on_log(self, logs, prefix=''): + n_steps = logs.pop('n_steps') + self._log_callback(logs, n_steps) + if prefix: + logs = {f'{prefix}{k}': v for k, v in logs.items()} + self._last_logged_metrics = dict(logs) + self.call_event('on_log', logs=logs) + + def _log_callback(self, logs, n_steps): + args = self.args + config = self.config + if config.num_moe_experts is not None: + moe_loss_scale = 1 / args.num_microbatches / n_steps + track_names = [] + load_balancing_type = config.moe_router_load_balancing_type + if isinstance(load_balancing_type, str): + load_balancing_type = [load_balancing_type] + if 'aux_loss' in load_balancing_type: + track_names.append('load_balancing_loss') + if 'seq_aux_loss' in load_balancing_type: + track_names.append('seq_load_balancing_loss') + if 'global_aux_loss' in load_balancing_type: + track_names.append('global_load_balancing_loss') + if config.moe_z_loss_coeff is not None: + track_names.append('z_loss') + track_moe_metrics( + loss_scale=moe_loss_scale, + iteration=self.state.iteration, + writer=None, + total_loss_dict=logs, + force_initialize=True, + track_names=track_names, + num_layers=config.num_layers, + moe_layer_freq=config.moe_layer_freq, + mtp_num_layers=args.mtp_num_layers) + if args.mtp_num_layers is not None: + mtp_loss_scale = 1 / args.num_microbatches / n_steps + mtp_logs = {} + MTPLossLoggingHelper.track_mtp_metrics(mtp_loss_scale, self.state.iteration, None, None, mtp_logs) + logs.update({k.replace(' ', '_'): v for k, v in mtp_logs.items()}) + # Track sparse attention indexer loss. + if args.dsa_indexer_loss_coeff is not None and args.dsa_indexer_loss_coeff > 0: + from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossLoggingHelper + indexer_loss_scale = 1 / args.num_microbatches / n_steps + idx_logs = {} + DSAIndexerLossLoggingHelper.track_indexer_metrics( + loss_scale=indexer_loss_scale, + iteration=self.state.iteration, + writer=None, + wandb_writer=None, + total_loss_dict=idx_logs, + ) + logs.update({k.replace(' ', '_'): v for k, v in idx_logs.items()}) + + for k, v in logs.items(): + if isinstance(v, torch.Tensor): + if v.numel() == 2: + v = v[0] / v[1] + v = v.item() + logs[k] = v + + def prepare_model(self): + args = self.args + self.unwrapped_models = get_mcore_model(args, self.template.config) + self.config = self.unwrapped_models[0].config + logger.info(f'model_config: {self.config}') + self.bridge = self.config.bridge + self.peft_models = self._prepare_peft_model(self.unwrapped_models) + self.wrapped_models = wrap_model(args, self.unwrapped_models) + + def _prepare_peft_model(self, models): + args = self.args + if args.mcore_model is None: + self.bridge.load_weights(models, args.model_dir) + peft_models = [prepare_mcore_model(args, model) for model in models] + if args.tuner_type == 'lora' and args.adapters and args.mcore_adapter is None: + assert len(args.adapters) == 1, 'Currently only support one adapter.' + self.bridge.load_weights(models, args.adapters[0], peft_format=True, adapter_name='default') + return peft_models + + def get_optimizer_and_scheduler(self): + args = self.args + if mcore_016: + from megatron.core.optimizer import AdamOptimizerConfig, SGDOptimizerConfig + if args.optimizer == 'adam' or 'muon' in args.optimizer: + # TODO(deyuf): Muon needs both adam + muon but get() only receive one config + # So for now we keep using adam config that's back compat with old way + config_cls = AdamOptimizerConfig + elif args.optimizer == 'sgd': + config_cls = SGDOptimizerConfig + else: + raise ValueError(f'Invalid optimizer type: {args.optimizer}') + else: + config_cls = OptimizerConfig + + kwargs = { + f.name: getattr(args, f.name) + for f in dataclasses.fields(config_cls) if hasattr(args, f.name) and f.name != 'loss_scale' + } + config = config_cls(**kwargs) + + if args.apply_wd_to_qk_layernorm or self.args.vit_lr is not None or self.args.aligner_lr is not None: + param_groups_context = self._patch_get_param_groups() + else: + param_groups_context = nullcontext() + with param_groups_context: + + if 'muon' not in config.optimizer: + # If the user is asking for a non-zero embedding init std, skip weight decay for embeddings + # to avoid embeddings from shrinking to zero as recommended in https://arxiv.org/abs/2312.16903 + # default_skip_embedding_weight_decay=args.embedding_init_method_std is not None, + optimizer = get_megatron_optimizer( + config, + self.wrapped_models, + ) + else: + from megatron.core.optimizer.muon import get_megatron_muon_optimizer + optimizer = get_megatron_muon_optimizer( + config, + self.wrapped_models, + layer_wise_distributed_optimizer='dist' in config.optimizer, + ) + opt_param_scheduler = get_optimizer_param_scheduler(args, optimizer) + return optimizer, opt_param_scheduler + + def _get_data_collator(self): + data_collator = self.template.data_collator + padding_to = get_padding_to(self.args) + logger.info(f'padding_to: {padding_to}') + data_collator = partial(data_collator, padding_to=padding_to) + return data_collator + + def cyclic_iter(self, iterable, use_origin_cyclic: bool = False): + training = self.unwrapped_models[0].training + if not training or use_origin_cyclic: + while True: + for x in iterable: + yield x + return + + args = self.args + epoch = 0 + is_finished = False + while True: + if not is_finished: + logger.info(f'The training of Epoch {epoch} starts...') + if hasattr(iterable, 'set_epoch'): + iterable.set_epoch(epoch) + for x in iterable: + yield x + # streaming + if training and args.num_train_epochs and epoch >= args.num_train_epochs - 1: + is_finished = True + epoch += 1 + if is_finished: + # Note that this approach will train for one additional step. + logger.info(f'Training of {epoch} epochs has been completed, the training has finished.') + args.train_iters = self.state.iteration + 1 + + def _get_param_groups_mcore_016( + self, + model_chunks: List[MegatronModule], + config: OptimizerConfig, + config_overrides, + ) -> List[Dict]: + return self._get_param_groups( + model_chunks, + no_weight_decay_cond=None, + scale_lr_cond=None, + lr_mult=1., + lr=config.lr, + min_lr=config.min_lr, + decoupled_lr=config.decoupled_lr, + decoupled_min_lr=config.decoupled_min_lr, + ) + + # Code borrowed from Megatron-LM + def _get_param_groups( + self, + model_chunks: List[MegatronModule], + no_weight_decay_cond: Optional[Callable], + scale_lr_cond: Optional[Callable], + lr_mult: float, + lr: float, + min_lr: float, + decoupled_lr: Optional[float], + decoupled_min_lr: Optional[float], + default_skip_embedding_weight_decay: bool = False, + ) -> List[Dict]: + """Create parameter groups for optimizer. + + Creates parameter groups based on weight decay condition (regularized vs + non regularized), learning rate scale condition (lr vs lr_mult * lr), + and whether it is expert parameters. scale_lr_cond is used during finetuning + where head of the network requires a scaled version of the base learning rate. + + Args: + model_chunks (List[MegatronModule]): model chunks to create parameter + groups for. + no_weight_decay_cond (func, optional): function to determine whether a + parameter should not perform weight decay. + scale_lr_cond (func, optional): function to determine whether a parameter + should have a scaled learning rate. + lr_mult (float): learning rate multiplier for parameters that + satisfy scale_lr_cond. + lr (float): learning rate. + min_lr (float): minimum learning rate. + decoupled_lr (Optional[float]): optional decoupled learning rate. + decoupled_min_lr (Optional[float]): optional decoupled minimum learning rate. + default_skip_embedding_weight_decay (bool): whether to skip weight decay for embedding + parameters by default, if no_weight_decay_cond is not provided. + + Returns: + List of parameter groups. + """ + args = self.args + if decoupled_min_lr is None: + decoupled_min_lr = min_lr + is_multimodal = args.megatron_model_meta.is_multimodal + if args.vit_lr is not None or args.aligner_lr is not None: + assert is_multimodal, 'vit_lr and aligner_lr are only supported for multimodal models.' + vit_lr = args.vit_lr if args.vit_lr is not None else args.lr + aligner_lr = args.aligner_lr if args.aligner_lr is not None else args.lr + logger.info_once(f'vit_lr: {vit_lr}, aligner_lr: {aligner_lr}, llm_lr: {args.lr}') + use_decoupled_learning_rate = decoupled_lr is not None + + # Map (wd_mult, lr_mult, is_expert_parallel, is_decoupled_lr) to params. + params_map = {} + for model_chunk in model_chunks: + visual = model_chunk.module.module.visual if is_multimodal else None + for name, param in model_chunk.named_parameters(): + if not param.requires_grad: + continue + + is_expert_parallel = not getattr(param, 'allreduce', True) + + if no_weight_decay_cond is not None: + no_wd: bool = no_weight_decay_cond(name, param) + elif args.apply_wd_to_qk_layernorm and any( + name.endswith(k) for k in ['q_layernorm.weight', 'k_layernorm.weight']): + no_wd = False + else: + # Do not regularize biases and norm parameters. + # optionally, also skip weight decay for embedding parameters if requested + # (useful if you do not want embeddings to shrink to zero in training + # https://arxiv.org/abs/2312.16903) + no_wd = ( + name.endswith('.bias') or len(param.shape) == 1 + or (default_skip_embedding_weight_decay and 'embedding' in name)) + _lr_mult = lr_mult + if scale_lr_cond is not None: + scale_lr = scale_lr_cond(name, param) + else: + scale_lr = False + # Handling multimodal models: vit_lr, aligner_lr + unwrapped_name = name.removeprefix('module.').removeprefix('module.') + if visual is not None: + is_aligner = any(unwrapped_name.startswith(f'visual.{k}') for k in visual._aligner or []) + is_vit = any(unwrapped_name.startswith(f'visual.{k}') + for k in visual._vision_tower) and not is_aligner + else: + is_aligner, is_vit = False, False + if is_vit and args.vit_lr: + scale_lr = True + _lr_mult = args.vit_lr / lr + elif is_aligner and args.aligner_lr: + scale_lr = True + _lr_mult = args.aligner_lr / lr + + if not no_wd and not scale_lr: + wd_mult, _lr_mult = 1.0, 1.0 + elif not no_wd and scale_lr: + wd_mult, _lr_mult = 1.0, _lr_mult + elif no_wd and not scale_lr: + wd_mult, _lr_mult = 0.0, 1.0 + else: + wd_mult, _lr_mult = 0.0, _lr_mult + + is_decoupled_lr = False + # For input/embedding and output layer: embedding.word_embeddings.weight / + # output_layer.weight. + if use_decoupled_learning_rate and getattr(param, 'is_embedding_or_output_parameter', False): + is_decoupled_lr = True + + key = (wd_mult, _lr_mult, is_expert_parallel, is_decoupled_lr) + if key not in params_map: + params_map[key] = [] + params_map[key].append(param) + + # Distributed checkpoint requires all ranks to have the same param groups, + # so we need to align the param groups across ranks, otherwise we may have + # runtime error when loading the checkpoint or numerical error when resuming training. + params_key = list(params_map.keys()) + gathered_params_key = [None for _ in range(torch.distributed.get_world_size())] + torch.distributed.all_gather_object(gathered_params_key, params_key) + for keys in gathered_params_key: + for key in keys: + if key not in params_key: + params_key.append(key) + + param_groups = [] + for key in params_key: + wd_mult, _lr_mult, is_expert_parallel, is_decoupled_lr = key + params = params_map[key] if key in params_map else [] + param_group = { + 'params': params, + 'wd_mult': wd_mult, + 'lr_mult': _lr_mult, + 'is_expert_parallel': is_expert_parallel, + 'is_decoupled_lr': is_decoupled_lr, + } + # Ensure param_group has required keys for matching when loading optimizer state + # See MegatronOptimizer._filter_and_reorder_param_groups. + if param_group_identifier_keys is not None: + assert set(param_group.keys()) - set(param_group_identifier_keys) == {'params'} + param_groups.append(param_group) + + # Update min and max lr in param groups + # These changes are compatible with mcore 0.16. + for param_group in param_groups: + if param_group['is_decoupled_lr']: + assert decoupled_lr is not None + param_group['max_lr'] = decoupled_lr + param_group['min_lr'] = decoupled_min_lr + else: + param_group['max_lr'] = lr + param_group['min_lr'] = min_lr + lr_mult = param_group.pop('lr_mult') + # Instead of using lr_mult to control the learning rate, we directly use max_lr/min_lr. + param_group['lr_mult'] = 1. + param_group['max_lr'] *= lr_mult + param_group['min_lr'] *= lr_mult + return param_groups + + @contextmanager + def _patch_get_param_groups(self): + from megatron.core import optimizer + _get_param_groups = optimizer._get_param_groups + if mcore_016: + optimizer._get_param_groups = self._get_param_groups_mcore_016 + else: + optimizer._get_param_groups = self._get_param_groups + try: + yield + finally: + optimizer._get_param_groups = _get_param_groups + + def _load_iteration(self): + args = self.args + ckpt_dir = None + if args.tuner_type == 'full': + ckpt_dir = args.model + else: + ckpt_dir = args.adapters[0] if args.adapters else args.model + if ckpt_dir is None: + return 0 + logger.info(f'checkpoint_dir: {ckpt_dir}') + tracker_path = os.path.join(ckpt_dir, 'latest_checkpointed_iteration.txt') + if not os.path.exists(tracker_path): + return 0 + with open(tracker_path, 'r') as f: + iteration = int(f.read()) + + common_path = os.path.join(ckpt_dir, f'iter_{iteration:07d}', 'common.pt') + if not os.path.exists(common_path): + return iteration + + state_dict = torch.load(common_path) + if 'args' not in state_dict: + return iteration + self.args.consumed_train_samples = getattr(state_dict['args'], 'consumed_train_samples', 0) + + return iteration + + def _prepare_vit_gradient_checkpointing(self, model): + visual = model.visual + if visual is None: + return + for vision_tower in visual._vision_tower: + module = deep_getattr(visual, vision_tower) + if self.args.vit_gradient_checkpointing: + dynamic_gradient_checkpointing(module, False) + try: + module.gradient_checkpointing_enable( + gradient_checkpointing_kwargs=self.args.vit_gradient_checkpointing_kwargs) + module.enable_input_require_grads() + except AttributeError: + pass + + @staticmethod + def _initialize_embedding(model): + # compat new_special_tokens + init_method = model.config.init_method + if hasattr(model, 'language_model'): + model = model.language_model + for key in ['embedding.word_embeddings', 'output_layer']: + if key == 'output_layer' and model.share_embeddings_and_output_weights: + continue + module = deep_getattr(model, key) + if module is None: + continue + initialize_mask = (module.weight == 0).all(dim=-1) + num_to_initialize = initialize_mask.sum().item() + if num_to_initialize == 0: + continue + logger.info_if(f'num_to_initialize: {num_to_initialize}', cond=mpu.get_data_parallel_rank() == 0) + tensor = module.weight.new_empty(num_to_initialize, module.weight.shape[1]) + module.weight.data[initialize_mask] = init_method(tensor) + if getattr(module.weight, 'main_param', None) is not None: + module.weight.main_param.copy_(module.weight.view(-1)) + + def _all_reduce_metric(self, + metric: Dict[str, torch.Tensor], + reduction=torch.distributed.ReduceOp.AVG, + group=None) -> Dict[str, torch.Tensor]: + if group is None: + group = mpu.get_data_parallel_group() + reporting_metric = torch.stack(list(metric.values()), dim=0) + torch.distributed.all_reduce(reporting_metric, reduction, group=group) + return {k: reporting_metric[i] for i, k in enumerate(metric.keys())} + + def merge_lora_adapters(self, adapter_name='default'): + """Merge LoRA adapters into base model weights for vLLM inference.""" + with torch.no_grad(): + for model in self.unwrapped_models: + for module in model.modules(): + if isinstance(module, LoraParallelLinear): + # Merge all active adapters + module.merge(adapter_names=[adapter_name]) + + def unmerge_lora_adapters(self): + """Unmerge LoRA adapters to restore training state.""" + with torch.no_grad(): + for model in self.unwrapped_models: + for module in model.modules(): + if isinstance(module, LoraParallelLinear): + # Unmerge to restore separate LoRA weights for training + module.unmerge() + + @staticmethod + def copy_path(src_path: str, tgt_path: str): + if not is_master(): + return + if not os.path.exists(src_path): + raise FileNotFoundError(f'Source path does not exist: {src_path}') + + if os.path.isfile(src_path): + os.makedirs(os.path.dirname(tgt_path), exist_ok=True) + shutil.copy(src_path, tgt_path) + elif os.path.isdir(src_path): + shutil.copytree(src_path, tgt_path, dirs_exist_ok=True) + else: + raise ValueError(f'Source path is neither a file nor a directory: {src_path}') + + def _prepare_data_iterator(self, train_dataset, val_dataset=None, use_origin_cyclic: bool = False): + train_dataloader, val_dataloader = self._prepare_dataloader(train_dataset, val_dataset) + train_data_iterator = iter(self.cyclic_iter(train_dataloader, use_origin_cyclic=use_origin_cyclic)) + val_data_iterator = None + if val_dataset is not None: + val_data_iterator = iter(self.cyclic_iter(val_dataloader, use_origin_cyclic=use_origin_cyclic)) + return train_data_iterator, val_data_iterator + + def setup_model_training(self): + args = self.args + config = self.config + + for m in self.wrapped_models: + m.train() + + if args.is_multimodal: + for m in self.unwrapped_models: + self._prepare_vit_gradient_checkpointing(m) + + config.grad_scale_func = self.optimizer.scale_loss + if isinstance(self.wrapped_models[0], (DDP, megatron_FSDP)) and args.overlap_grad_reduce: + assert config.no_sync_func is None, ('When overlap_grad_reduce is True, config.no_sync_func must be None; ' + 'a custom no_sync_func is not supported when overlapping grad-reduce') + config.no_sync_func = [model_chunk.no_sync for model_chunk in self.wrapped_models] + if len(self.wrapped_models) == 1: + config.no_sync_func = config.no_sync_func[0] + if args.align_grad_reduce: + config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in self.wrapped_models] + if len(self.wrapped_models) == 1: + config.grad_sync_func = config.grad_sync_func[0] + if args.overlap_param_gather and args.align_param_gather: + config.param_sync_func = [model_chunk.start_param_sync for model_chunk in self.wrapped_models] + if len(self.wrapped_models) == 1: + config.param_sync_func = config.param_sync_func[0] + config.finalize_model_grads_func = finalize_model_grads + + self._start_iteration = self.state.iteration + self._pre_hook_enabled = False + if should_disable_forward_pre_hook(args): + disable_forward_pre_hook(self.wrapped_models, param_sync=False) + self._saved_param_sync_func = config.param_sync_func + config.param_sync_func = None + + if args.nccl_comm_warmup: + # Eagerly create NCCL communicators while GPU memory is still free. Lazily-initialized + # comms (e.g. the dp/cp loss all-reduce and grad-sync coalescing) otherwise first fire + # at the iteration-1 memory peak, where NCCL's internal cudaMalloc can fail with + # "Failed to CUDA calloc async N bytes". A 1-element dummy all-reduce per group is + # numerically inert and forces the communicator to be created up front. + dummy = torch.zeros(1, device=get_current_device()) + warmed = 0 + for getter, kwargs in ( + (mpu.get_data_parallel_group, { + 'with_context_parallel': True + }), + (mpu.get_data_parallel_group, {}), + (mpu.get_context_parallel_group, {}), + (mpu.get_tensor_model_parallel_group, {}), + (mpu.get_pipeline_model_parallel_group, {}), + (mpu.get_model_parallel_group, {}), + (mpu.get_embedding_group, {}), + (mpu.get_position_embedding_group, {}), + ): + try: + group = getter(**kwargs) + except (AssertionError, ValueError, TypeError): + continue + for g in (group if isinstance(group, list) else [group]): + if g is not None: + torch.distributed.all_reduce(dummy, group=g) + warmed += 1 + torch.cuda.synchronize() + logger.info(f'NCCL communicator warm-up done ({warmed} groups).') + + self.call_event('on_train_begin') + self._train_metrics = {} + + def setup_training(self, train_dataset, val_dataset): + self.setup_model_training() + args = self.args + if args.virtual_pipeline_model_parallel_size is not None: + train_data_iterator, val_data_iterator = [], [] + for _ in range(args.virtual_pipeline_model_parallel_size): + train_it, val_it = self._prepare_data_iterator(train_dataset, val_dataset) + train_data_iterator.append(train_it) + val_data_iterator.append(val_it) + else: + train_data_iterator, val_data_iterator = self._prepare_data_iterator(train_dataset, val_dataset) + return train_data_iterator, val_data_iterator + + def run_train_step(self, train_data_iterator, val_data_iterator): + """Execute one training iteration including logging / eval / save. + + Returns ``True`` if training should continue, ``False`` if done. + """ + args = self.args + config = self.config + state = self.state + + self.call_event('on_step_begin') + maybe_finalize_async_save(args, blocking=False) + metrics, grad_norm, update_successful = self.train_step(train_data_iterator) + + if state.iteration == self._start_iteration: + if update_successful: + if should_disable_forward_pre_hook(args): + enable_forward_pre_hook(self.wrapped_models) + config.param_sync_func = self._saved_param_sync_func + self._pre_hook_enabled = True + else: + self._start_iteration = state.iteration + 1 + + state.iteration += 1 + self.call_event('on_step_end') + self._aggregated_metrics(metrics, self._train_metrics) + self._train_metrics['grad_norm'] = grad_norm + for param_group in self.optimizer.param_groups: + if len(param_group['params']) == 0: + continue + self._train_metrics['learning_rate'] = param_group['lr'] + break + if state.should_log: + state.should_log = False + self.on_log(logs=self._train_metrics) + self._train_metrics = {} + + eval_metrics = None + if state.should_eval: + state.should_eval = False + if should_disable_forward_pre_hook(args): + disable_forward_pre_hook(self.wrapped_models) + self._pre_hook_enabled = False + eval_metrics = self.evaluate(val_data_iterator) + for m in self.wrapped_models: + m.train() + if should_disable_forward_pre_hook(args): + enable_forward_pre_hook(self.wrapped_models) + self._pre_hook_enabled = True + + if state.should_save: + self._determine_best_metric(eval_metrics) + if should_disable_forward_pre_hook(args): + disable_forward_pre_hook(self.wrapped_models) + state.should_save = False + self.save_checkpoint() + self.call_event('on_save', output_dir=self.state.last_model_checkpoint) + if should_disable_forward_pre_hook(args): + enable_forward_pre_hook(self.wrapped_models) + + return state.iteration < args.train_iters + + def finalize_training(self): + """Cleanup after training loop completes.""" + self.call_event('on_train_end') + if self._pre_hook_enabled: + disable_forward_pre_hook(self.wrapped_models) + maybe_finalize_async_save(self.args, blocking=True, terminate=True) + + def train(self, train_dataset, val_dataset): + train_data_iterator, val_data_iterator = self.setup_training(train_dataset, val_dataset) + while self.state.iteration < self.args.train_iters: + self.run_train_step(train_data_iterator, val_data_iterator) + self.finalize_training() + + def _determine_best_metric(self, metrics) -> bool: + args = self.args + state = self.state + if (args.metric_for_best_model is None or metrics is None or not is_last_rank() + or args.metric_for_best_model not in metrics): + if args.metric_for_best_model is None: + return False + tensor = torch.zeros((3, ), device=get_current_device(), dtype=torch.float64) + torch.distributed.all_reduce(tensor) + is_new_best_metric = bool(tensor[2].item()) + if is_new_best_metric: + state.best_metric = tensor[0].item() + state.best_global_step = int(tensor[1].item()) + return is_new_best_metric + metric_value = metrics[args.metric_for_best_model] + op = operator.ge if args.greater_is_better else operator.le + if state.best_metric is None: + state.best_metric = float('-inf') if args.greater_is_better else float('inf') + + is_new_best_metric = False + if op(metric_value, state.best_metric): + state.best_metric = metric_value + state.best_global_step = state.global_step + is_new_best_metric = True + tensor = torch.tensor([state.best_metric, state.best_global_step, is_new_best_metric], + device=get_current_device(), + dtype=torch.float64) + torch.distributed.all_reduce(tensor) + return is_new_best_metric + + def save_checkpoint(self): + args = self.args + state = self.state + args.consumed_train_samples = state.consumed_train_samples + iteration = state.iteration + output_dir = os.path.join(args.output_dir, f'checkpoint-{iteration}') + os.makedirs(output_dir, exist_ok=True) + args_path = os.path.join(os.path.dirname(output_dir), 'args.json') + self.copy_path(args_path, os.path.join(output_dir, 'args.json')) + if args.save_safetensors and args.no_save_optim: + model = [] + else: + model = self.wrapped_models + gc_collect() + save_mcore_checkpoint( + args, + model, + self.optimizer, + self.opt_param_scheduler, + iteration=iteration, + peft_format=args.tuner_type == 'lora', + output_dir=output_dir) + state.last_model_checkpoint = output_dir + if state.best_global_step is not None: + best_model_checkpoint = os.path.join(args.output_dir, f'checkpoint-{state.best_global_step}') + if os.path.exists(best_model_checkpoint): + state.best_model_checkpoint = best_model_checkpoint + # safetensors + if args.save_safetensors: + skip_saving_adapter = args.tuner_type == 'lora_llm' or ( + args.tuner_type == 'lora' and args.merge_lora and not hasattr(self.bridge, '_support_hf_grouped_lora')) + save_missing_weights = args.save_missing_weights and args.model_dir + + if not skip_saving_adapter: + self.bridge.save_weights( + self.unwrapped_models, + output_dir, + peft_format=args.tuner_type == 'lora', + args=args, + processor=self.template.processor, + save_missing_weights=save_missing_weights, + ) + # merge-lora does not store lora, lora saving may report an error (Qwen3-VL-Moe) + if args.tuner_type != 'full' and args.merge_lora: + self.merge_lora_adapters() + origin_output_dir = output_dir + output_dir = f'{output_dir}-merged' + os.makedirs(output_dir, exist_ok=True) + for fname in ['latest_checkpointed_iteration.txt', 'args.json']: + src_path = os.path.join(origin_output_dir, fname) + self.copy_path(src_path, os.path.join(output_dir, fname)) + # common.pt + common_path = os.path.join(origin_output_dir, f'iter_{iteration:07d}', 'common.pt') + if os.path.exists(common_path): + tgt_common_path = os.path.join(output_dir, f'iter_{iteration:07d}', 'common.pt') + os.makedirs(os.path.dirname(tgt_common_path), exist_ok=True) + self.copy_path(common_path, tgt_common_path) + self.bridge.save_weights( + self.unwrapped_models, + output_dir, + peft_format=False, + args=args, + processor=self.template.processor, + save_missing_weights=save_missing_weights, + ) + self.unmerge_lora_adapters() + + if is_master(): + self._rotate_checkpoints(args.output_dir) + + def _rotate_checkpoints(self, output_dir: str): + # Code borrowed from huggingface/transformers + args = self.args + if args.save_total_limit is None or args.save_total_limit <= 0: + return + + checkpoints_sorted = self._sorted_checkpoints(output_dir) + if len(checkpoints_sorted) <= args.save_total_limit: + return + + number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - args.save_total_limit) + checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete] + for checkpoint in checkpoints_to_be_deleted: + shutil.rmtree(checkpoint, ignore_errors=True) + if os.path.exists(f'{checkpoint}-merged'): + shutil.rmtree(f'{checkpoint}-merged', ignore_errors=True) + + def _sorted_checkpoints(self, output_dir: str): + # Code borrowed from huggingface/transformers + state = self.state + glob_checkpoints = [ + str(p) for p in Path(output_dir).glob('checkpoint-*') if p.is_dir() and not p.name.endswith('-merged') + ] + # Sort by modification time + checkpoints_sorted = sorted(glob_checkpoints, key=os.path.getmtime) + + # Make sure we don't delete the best model. + if state.best_model_checkpoint is not None and state.best_model_checkpoint in checkpoints_sorted: + best_model_index = checkpoints_sorted.index(state.best_model_checkpoint) + checkpoints_sorted.pop(best_model_index) + checkpoints_sorted.append(state.best_model_checkpoint) + return checkpoints_sorted + + def training_log(self, metrics, grad_norm): + learning_rate = None + for param_group in self.optimizer.param_groups: + if len(param_group['params']) == 0: + continue + learning_rate = param_group['lr'] + logger.info(f'metrics: {metrics}, grad_norm: {grad_norm}, learning_rate: {learning_rate}') + + def evaluate(self, val_data_iterator): + args = self.args + for m in self.wrapped_models: + m.eval() + eval_metrics = {} + forward_backward_func = get_forward_backward_func() + self.call_event('on_eval_begin') + with torch.no_grad(): + for _ in range(args.eval_iters): + data_iterator = self._replace_data_iterator(val_data_iterator) + metrics = forward_backward_func( + forward_step_func=self.forward_step, + data_iterator=data_iterator, + model=self.wrapped_models, + num_microbatches=self.args.num_microbatches, + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + forward_only=True, + ) + self.call_event('on_eval_step') + self._aggregated_metrics(metrics, eval_metrics) + self.compute_eval_metrics(eval_metrics) + self.on_log(logs=eval_metrics, prefix='eval_') + self.call_event('on_eval_end') + return eval_metrics + + def compute_eval_metrics(self, metrics): + if self.eval_metrics is not None and mpu.is_pipeline_last_stage(): + metric = self.eval_metrics.compute() + for k, v in metric.items(): + metrics[k] = v if isinstance(v, torch.Tensor) else torch.tensor(v) + self.eval_metrics.reset() + + def _replace_data_iterator(self, data_iterator): + return data_iterator + + def train_step(self, train_data_iterator): + args = self.args + forward_backward_func = get_forward_backward_func() + for m in self.wrapped_models: + m.zero_grad_buffer() + self.optimizer.zero_grad() + # TODO: refactor _replace_data_iterator + data_iterator = self._replace_data_iterator(train_data_iterator) + + if self.enable_routing_replay: + RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + + metrics = forward_backward_func( + forward_step_func=self.forward_step, + data_iterator=data_iterator, + model=self.wrapped_models, + num_microbatches=args.num_microbatches, + seq_length=args.seq_length, + micro_batch_size=args.micro_batch_size, + forward_only=False, + ) + + update_successful, grad_norm, _ = self.optimizer.step() + update_successful = logical_and_across_model_parallel_group(update_successful) + grad_norm = reduce_max_stat_across_model_parallel_group(grad_norm) + if update_successful: + self.opt_param_scheduler.step(increment=args.global_batch_size) + + if self.enable_routing_replay: + RouterReplay.clear_global_router_replay_action() + RouterReplay.clear_global_indices() + + return metrics, grad_norm, update_successful + + def _aggregated_metrics(self, metrics, total_metrics): + if 'n_steps' not in total_metrics: + total_metrics['n_steps'] = 0 + total_metrics['n_steps'] += 1 + if not metrics: + return + metrics = RowPreprocessor.rows_to_batched(metrics) + for key, val in metrics.items(): + val = torch.stack([v for v in val if v is not None], dim=0) + if val[0].numel() == 2: + val = val.sum(dim=0) + if val[1] == 0: + continue + elif val[0].numel() == 1: + val = val.new_tensor([val.sum(), val.shape[0]]) + else: + raise ValueError(f'Invalid value shape: {val[0].shape} for key {key}') + if key not in total_metrics: + total_metrics[key] = torch.tensor([0.0, 0.0], dtype=torch.float32, device=torch.cuda.current_device()) + total_metrics[key] += val + + def _prepare_dataloader(self, train_dataset, val_dataset=None): + args = self.args + val_dataloader = None + if args.streaming: + train_dataloader = build_streaming_dataloader(args, train_dataset, self.data_collator) + if val_dataset is not None: + val_dataloader = build_streaming_dataloader(args, val_dataset, self.data_collator) + return train_dataloader, val_dataloader + train_batch_sampler = MegatronPretrainingRandomSampler( + train_dataset, + total_samples=len(train_dataset), + consumed_samples=self.state.consumed_train_samples, + micro_batch_size=args.micro_batch_size, + data_parallel_rank=mpu.get_data_parallel_rank(), + data_parallel_size=mpu.get_data_parallel_world_size(), + data_sharding=args.data_sharding, + shuffle=args.train_dataloader_shuffle, + group_by_length=args.group_by_length, + ) + train_dataloader = self._create_dataloader(train_dataset, train_batch_sampler) + if val_dataset is not None: + val_batch_sampler = MegatronPretrainingSampler( + total_samples=len(val_dataset), + consumed_samples=0, + micro_batch_size=args.micro_batch_size, + data_parallel_rank=mpu.get_data_parallel_rank(), + data_parallel_size=mpu.get_data_parallel_world_size(), + ) + val_dataloader = self._create_dataloader(val_dataset, val_batch_sampler) + return train_dataloader, val_dataloader + + def _create_dataloader(self, dataset, batch_sampler): + args = self.args + + dataloader = torch.utils.data.DataLoader( + dataset, + batch_sampler=batch_sampler, + num_workers=args.dataloader_num_workers, + pin_memory=args.dataloader_pin_memory, + persistent_workers=args.dataloader_persistent_workers if args.dataloader_num_workers > 0 else False, + prefetch_factor=args.dataloader_prefetch_factor if args.dataloader_num_workers > 0 else None, + collate_fn=self.data_collator, + ) + return dataloader + + @abstractmethod + def forward_step(self, data_iterator, model): + pass + + def _prepare_batch(self, data, vp_stage=None): + return prepare_batch(self.args, data, vp_stage=vp_stage) + + def get_batch(self, data_iterator, vp_stage=None): + """Generate a batch.""" + return self._prepare_batch(next(data_iterator), vp_stage) + + def _collect_config_info(self) -> Dict[str, str]: + """ + Collects trainer-specific configuration details. + + Subclasses can override this method to provide additional configuration + information for model compatibility verification. + + Returns: + Dict[str, str]: Configuration parameters as key-value pairs. + """ + if self.__class__.__name__ == 'MegatronTrainer': + if not self.template.use_chat_template: + return { + 'seq2seq_mode': 'pt', + } + else: + return { + 'seq2seq_mode': 'sft', + } + return {} + + def get_last_tokens(self, output_tensor, packed_seq_params=None, attention_mask=None): + if self.args.context_parallel_size > 1: + output_tensor = reconstruct_tensor_cp(output_tensor, packed_seq_params, dim=1) + if packed_seq_params is None: + # Compatible with attention_mask_2d + if attention_mask.dim() > 2: + attention_mask = (~attention_mask).sum(dim=(1, 2)) > 0 + last_token_idx = get_last_valid_indices(attention_mask.long()) + last_tokens = output_tensor[torch.arange(output_tensor.shape[0]), last_token_idx] + else: + num_samples = packed_seq_params.seq_lens.shape[0] + last_token_idx = packed_seq_params.cu_seqlens_q[:num_samples] + packed_seq_params.seq_lens - 1 + last_tokens = output_tensor[0, last_token_idx] + return last_tokens diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/batch_sampler.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/batch_sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..dfbc788bbe5b758c3f63dbf0b68d65c327ec87ee --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/batch_sampler.py @@ -0,0 +1,155 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch + +from swift.utils import get_logger + +logger = get_logger() + + +# Code borrowed from megatron-lm +class MegatronPretrainingSampler: + + def __init__(self, + total_samples, + consumed_samples, + micro_batch_size, + data_parallel_rank, + data_parallel_size, + drop_last=True): + # Keep a copy of input params for later use. + self.total_samples = total_samples + self.consumed_samples = consumed_samples + self.micro_batch_size = micro_batch_size + self.data_parallel_rank = data_parallel_rank + self.micro_batch_times_data_parallel_size = \ + self.micro_batch_size * data_parallel_size + self.drop_last = drop_last + + # Sanity checks. + assert self.total_samples > 0, \ + 'no sample to consume: {}'.format(self.total_samples) + assert self.consumed_samples < self.total_samples, \ + 'no samples left to consume: {}, {}'.format(self.consumed_samples, + self.total_samples) + assert self.micro_batch_size > 0 + assert data_parallel_size > 0 + assert self.data_parallel_rank < data_parallel_size, \ + 'data_parallel_rank should be smaller than data size: {}, ' \ + '{}'.format(self.data_parallel_rank, data_parallel_size) + + def __len__(self): + return self.total_samples + + def get_start_end_idx(self): + start_idx = self.data_parallel_rank * self.micro_batch_size + end_idx = start_idx + self.micro_batch_size + return start_idx, end_idx + + def __iter__(self): + batch = [] + # Last batch will be dropped if drop_last is not set False + for idx in range(self.consumed_samples, self.total_samples): + batch.append(idx) + if len(batch) == self.micro_batch_times_data_parallel_size: + start_idx, end_idx = self.get_start_end_idx() + yield batch[start_idx:end_idx] + batch = [] + + # Check the last partial batch and see drop_last is set + if len(batch) > 0 and not self.drop_last: + start_idx, end_idx = self.get_start_end_idx() + yield batch[start_idx:end_idx] + + +# Code borrowed from megatron-lm +class MegatronPretrainingRandomSampler: + + def __init__( + self, + dataset, + total_samples, + consumed_samples, + micro_batch_size, + data_parallel_rank, + data_parallel_size, + data_sharding, + shuffle: bool = True, + group_by_length: bool = False, + ): + # Keep a copy of input params for later use. + self.dataset = dataset + self.total_samples = total_samples + self.consumed_samples = consumed_samples + self.micro_batch_size = micro_batch_size + self.data_parallel_rank = data_parallel_rank + self.data_parallel_size = data_parallel_size + if group_by_length: + if data_sharding: + data_sharding = False + logger.warning('`group_by_length=True` is incompatible with `data_sharding=True`. ' + 'Setting `data_sharding=False` to enable length grouping.') + if not shuffle: + raise ValueError('shuffle must be True when group_by_length is True') + self.data_sharding = data_sharding + self.shuffle = shuffle + self.group_by_length = group_by_length + self.lengths = self.dataset['lengths'] if group_by_length else None + if self.lengths is not None: + self.lengths = [max(length) if isinstance(length, list) else length for length in self.lengths] + self.micro_batch_times_data_parallel_size = self.micro_batch_size * data_parallel_size + self.last_batch_size = self.total_samples % self.micro_batch_times_data_parallel_size + + # Sanity checks. + assert self.total_samples > 0, 'no sample to consume: {}'.format(self.total_samples) + assert self.micro_batch_size > 0 + assert data_parallel_size > 0 + assert self.data_parallel_rank < data_parallel_size, ( + 'data_parallel_rank should be smaller than data size: {}, ' + '{}'.format(self.data_parallel_rank, data_parallel_size)) + + def __len__(self): + return self.total_samples + + def __iter__(self): + active_total_samples = self.total_samples - self.last_batch_size + self.epoch = self.consumed_samples // active_total_samples + current_epoch_samples = self.consumed_samples % active_total_samples + assert current_epoch_samples % self.micro_batch_times_data_parallel_size == 0 + + if self.shuffle: + # data sharding and random sampling + if self.data_sharding: + bucket_size = (self.total_samples // self.micro_batch_times_data_parallel_size) * self.micro_batch_size + bucket_offset = current_epoch_samples // self.data_parallel_size + start_idx = self.data_parallel_rank * bucket_size + + g = torch.Generator() + g.manual_seed(self.epoch) + random_idx = torch.randperm(bucket_size, generator=g).tolist() + idx_range = [start_idx + x for x in random_idx[bucket_offset:]] + else: + full_bucket_size = (self.total_samples // self.micro_batch_size) * self.micro_batch_size + full_bucket_offset = current_epoch_samples + g = torch.Generator() + g.manual_seed(self.epoch) + if self.group_by_length: + from transformers.trainer_pt_utils import get_length_grouped_indices + idx_range_total = get_length_grouped_indices( + self.lengths, self.micro_batch_times_data_parallel_size, generator=g) + else: + idx_range_total = torch.randperm(full_bucket_size, generator=g).tolist() + idx_range_active = idx_range_total[full_bucket_offset:] + idx_range = idx_range_active[self.data_parallel_rank::self.data_parallel_size] + else: + full_bucket_size = (self.total_samples // self.micro_batch_size) * self.micro_batch_size + full_bucket_offset = current_epoch_samples + idx_range = range(full_bucket_offset + self.data_parallel_rank, full_bucket_size, self.data_parallel_size) + + batch = [] + # Last batch if not complete will be dropped. + for idx in idx_range: + batch.append(idx) + if len(batch) == self.micro_batch_size: + self.consumed_samples += self.micro_batch_times_data_parallel_size + yield batch + batch = [] diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/dpo_trainer.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/dpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..98584b3c16b1617b118bfe02a8c319daa16483df --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/dpo_trainer.py @@ -0,0 +1,94 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from collections import namedtuple +from functools import partial +from megatron.core import mpu +from torch.distributed.nn import all_reduce + +from swift.rlhf_trainers import DPOTrainer +from swift.utils import get_current_device, get_logger +from .rlhf_mixin import MegatronRLHFTrainer + +logger = get_logger() + + +class DummyDPOTrainer(DPOTrainer): + # For reusing the dpo_loss function implemented in Swift's DPOTrainer. + def __init__(self, args): + self.accelerator = namedtuple('Accelerator', ['device'])(device=get_current_device()) + self.f_alpha_divergence_coef = 1. + self.f_divergence_params = {'alpha_divergence_coef': self.f_alpha_divergence_coef} + self.reference_free = args.reference_free + self.label_smoothing = args.label_smoothing + self.f_divergence_type = args.f_divergence_type + self.loss_type = args.loss_type + self.beta = args.beta + + +class MegatronDPOTrainer(MegatronRLHFTrainer): + + def __init__(self, args, template): + super().__init__(args, template) + self.dummy_dpo_trainer = DummyDPOTrainer(args) + + def loss_func(self, output_tensor: torch.Tensor, *, labels: torch.Tensor, packed_seq_params): + ref_output_tensor = output_tensor[:output_tensor.shape[0] // 2].detach() + output_tensor = output_tensor[output_tensor.shape[0] // 2:] + args = self.args + num_samples = labels.shape[0] if packed_seq_params is None else packed_seq_params.seq_lens.shape[0] + + logps = self.get_logps(output_tensor, labels, packed_seq_params) + ref_logps = self.get_logps(ref_output_tensor, labels, packed_seq_params) + loss, chosen_rewards, rejected_rewards = self.dummy_dpo_trainer.dpo_loss( + logps[:num_samples // 2], + logps[num_samples // 2:], + ref_logps[:num_samples // 2], + ref_logps[num_samples // 2:], + ) + if args.rpo_alpha: + loss_mask = labels != -100 + if args.padding_free: + num_tokens = packed_seq_params.cu_seqlens_q[num_samples // 2] // args.context_parallel_size + loss_mask[:, num_tokens:] = 0 + else: + loss_mask[num_samples // 2:] = 0 + nll_loss = torch.concat([torch.sum(output_tensor * loss_mask)[None], loss_mask.sum()[None]]) + if args.context_parallel_size > 1: + nll_loss = all_reduce(nll_loss, group=mpu.get_context_parallel_group()) + nll_loss = nll_loss[0] / nll_loss[1] + loss = loss + args.rpo_alpha * nll_loss + loss = loss.mean() + metric = { + 'loss': loss.detach().clone(), + 'logps/chosen': logps[:num_samples // 2].mean(), + 'logps/rejected': logps[num_samples // 2:].mean(), + 'rewards/chosen': chosen_rewards.mean(), + 'rewards/rejected': rejected_rewards.mean(), + 'rewards/accuracies': (chosen_rewards > rejected_rewards).float().mean(), + 'rewards/margins': (chosen_rewards - rejected_rewards).mean(), + } + if args.rpo_alpha: + metric['nll_loss'] = nll_loss.detach() + metric = self._all_reduce_metric(metric) + # fix megatron-lm bug + loss = loss / mpu.get_context_parallel_world_size() + return loss, metric + + def forward_step(self, data_iterator, model): + unwrapped_model = model.module.module + input_tensor = unwrapped_model.get_input_tensor() + vp_stage = unwrapped_model.vp_stage + data = self.get_batch(data_iterator, vp_stage) + data.pop('loss_scale', None) + # ref_model + with torch.no_grad(), self.null_ref_context() as ref_models: + ref_model = ref_models[vp_stage or 0] + if input_tensor is not None: + ref_model.set_input_tensor(input_tensor[:input_tensor.shape[0] // 2].detach()) + ref_output_tensor = ref_model(**data) + + if input_tensor is not None: + unwrapped_model.set_input_tensor(input_tensor[input_tensor.shape[0] // 2:]) + output_tensor = model(**data) + return torch.concat([ref_output_tensor, output_tensor], dim=0), partial( + self.loss_func, labels=data.get('labels'), packed_seq_params=data.get('packed_seq_params')) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/embedding_trainer.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/embedding_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..079b770f39c963832a067bfda4326c3eacccd6ec --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/embedding_trainer.py @@ -0,0 +1,62 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch.nn +import torch.nn.functional as F +from functools import partial + +from swift.loss import loss_map +from swift.metrics import eval_metrics_map +from swift.utils import get_logger +from .base import BaseMegatronTrainer + +logger = get_logger() + + +class MegatronEmbeddingTrainer(BaseMegatronTrainer): + + def __init__(self, args, template): + super().__init__(args, template) + self._loss_func = loss_map[args.loss_type](args, self) + eval_metric = 'infonce' if args.loss_type == 'infonce' else 'paired' + self.eval_metrics = eval_metrics_map[eval_metric](args, self) + + def loss_func(self, + output_tensor: torch.Tensor, + *, + labels: torch.Tensor, + packed_seq_params=None, + attention_mask=None): + training = self.unwrapped_models[0].training + last_hidden_state = self.get_last_tokens(output_tensor, packed_seq_params, attention_mask) + if not training: + self.eval_metrics.update(last_hidden_state.detach(), labels) + mrl_dims = self.args.mrl_dims + if mrl_dims: + # Matryoshka Representation Learning: compute loss on each truncated dimension + # and aggregate with the corresponding weights. + loss = None + for dim, weight in mrl_dims.items(): + if dim > last_hidden_state.shape[-1]: + logger.warning_once(f'MRL: skipping dimension {dim} because it exceeds the model hidden size ' + f'({last_hidden_state.shape[-1]}).') + continue + sliced = F.normalize(last_hidden_state[..., :dim], p=2, dim=-1) + cur_loss = weight * self._loss_func({'last_hidden_state': sliced}, labels) + loss = cur_loss if loss is None else loss + cur_loss + else: + loss = self._loss_func({'last_hidden_state': last_hidden_state}, labels) + metric = {'loss': loss.detach().clone()} + metric = self._all_reduce_metric(metric) + return loss, metric + + def forward_step(self, data_iterator, model): + vp_stage = model.module.module.vp_stage + data = self.get_batch(data_iterator, vp_stage) + labels = data.pop('labels', None) + output_tensor = model(**data) + loss_func = partial( + self.loss_func, + labels=labels, + packed_seq_params=data.get('packed_seq_params'), + attention_mask=data.get('attention_mask') + if data.get('attention_mask') is not None else data.get('attention_mask_2d')) + return output_tensor, loss_func diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/gkd_trainer.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/gkd_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..1d558ed76f1c3ac02371b23e5addcfd8031f7daf --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/gkd_trainer.py @@ -0,0 +1,427 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +import random +import torch +import torch.nn.functional as F +from contextlib import contextmanager +from functools import partial +from mcore_bridge import set_random_seed +from megatron.core import mpu +from transformers.utils import ContextManagers +from typing import Dict, List, Optional + +from swift.megatron.arguments import MegatronArguments +from swift.rl_core.data import GKDSample +from swift.rl_core.resample import resample_encode_failed_inputs +from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_output, build_opsd_samples, build_teacher_requests, + encode_gkd_samples, fetch_teacher_parsed_by_routing) +from swift.rlhf_trainers.gkd_loss import DataSource, TeacherOutput, gkd_loss +from swift.template import Template +from swift.utils import get_logger, to_device +from ..utils import forward_step_helper, get_padding_to +from .gkd_utils import cp_reduce, tp_gather_topk, vocab_parallel_topk +from .rlhf_mixin import MegatronRLHFTrainer +from .rollout_mixin import MegatronRolloutMixin +from .utils import gather_object +from .vocab_parallel_utils import vocab_parallel_kl_div, vocab_parallel_log_softmax + +logger = get_logger() + + +class MegatronGKDTrainer(MegatronRolloutMixin, MegatronRLHFTrainer): + + sample_cls = GKDSample + + def __init__(self, args: MegatronArguments, template, **kwargs): + self.vllm_client = kwargs.pop('vllm_client', None) + + # GKD-specific parameters + self.beta = args.beta # JSD interpolation coefficient + self.temperature = args.temperature + self.lmbda = args.lmbda # On-policy probability + self.args = args + self._setup_teacher() + if self._teacher_use_disable_adapter: + logger.info('Self-distillation mode: using disable_adapter() for fixed teacher (no extra model)') + self.sft_alpha = getattr(args, 'sft_alpha', 0.0) # Weight for SFT loss + + # GKD top-k logits configuration + self.gkd_logits_topk = getattr(args, 'gkd_logits_topk', None) + + self.use_vllm = getattr(args, 'use_vllm', False) + self.steps_per_generation = args.steps_per_generation + self.generation_batch_size = args.generation_batch_size + super().__init__(args, template) + + if self.use_teacher_api: + logger.info(f'Using teacher model API for logprobs, top_logprobs={self.gkd_logits_topk}') + + # Get device for data processing + self.device = torch.cuda.current_device() + + # Initialize vLLM rollout engine if on-policy generation is enabled + self._init_rollout_engine() + + # Truncation strategy for handling sequences that exceed max_length + self.truncation_strategy = args.truncation_strategy + self.max_completion_length = args.max_completion_length + + self.resample_data_iterator = None + self._buffered_inputs = None + + self._prepare_logging() + + def train(self, train_dataset, val_dataset): + if self.truncation_strategy == 'delete': + self.resample_data_iterator = self._init_resample_data_iterator(train_dataset) + super().train(train_dataset, val_dataset) + + def prepare_model(self): + super().prepare_model() + if self.use_teacher_api: + logger.info('Skipping local teacher model loading - using external API for teacher logprobs') + elif self._is_self_distillation: + logger.info('Self-distillation mode: using student model as teacher (no separate teacher loaded)') + self._load_teacher_model() + + @contextmanager + def _template_context(self, template: Template, max_length: Optional[int] = None): + """Context manager to temporarily modify max_length constraint from template.""" + original_max_length = template.max_length + template.max_length = max_length + try: + yield + finally: + template.max_length = original_max_length + + def _build_teacher_requests(self, samples: List[GKDSample]): + if not self.use_teacher_api: + return [] + return build_teacher_requests(samples, self.template) + + def _encode_samples(self, samples: List[GKDSample]) -> Dict[str, torch.Tensor]: + template = self.template + args = self.args + + with self._template_context(template): + student_encoded_list, teacher_encoded_list, has_opsd = encode_gkd_samples(samples, template) + + padding_to = get_padding_to(args) + encoded_batch = to_device(template.data_collator(student_encoded_list, padding_to=padding_to), self.device) + if has_opsd: + teacher_model_inputs = to_device( + template.data_collator(teacher_encoded_list, padding_to=padding_to), self.device) + else: + teacher_model_inputs = encoded_batch.copy() + encoded_batch['teacher_model_inputs'] = teacher_model_inputs + return encoded_batch + + def _get_random_num(self) -> float: + """Generate a deterministic random number consistent across all processes. + + Uses an isolated Random instance with seed based on args.seed + step counter + + Returns: + float: A random number in the range [0.0, 1.0). + """ + seed = int(getattr(self.args, 'seed', 0)) + seed += int(self._step) + rng = random.Random(seed) + return rng.random() + + def _determine_data_source(self) -> DataSource: + """Determine data source for current step based on GKD algorithm. + + GKD training mode selection logic: + 1. With probability lmbda: On-Policy (student generates) + 2. Otherwise: Off-Policy (use dataset responses) + + Returns: + DataSource enum indicating which source to use. + """ + random_num = self._get_random_num() + + if random_num < self.lmbda: + # Mode 1: On-Policy learning, student model generates responses + if self.use_vllm: + return DataSource.STUDENT + else: + # If vLLM not enabled, fall back to dataset + logger.warning_once('On-policy mode triggered but use_vllm=False. ' + 'Falling back to dataset responses. Enable vLLM for on-policy generation.') + return DataSource.DATASET + else: + # Mode 2: Off-Policy learning, use dataset responses + return DataSource.DATASET + + def _init_resample_data_iterator(self, train_dataset): + """Initialize an independent data iterator for resampling. + + Uses a different seed (args.seed + 1) to avoid overlapping with training samples. + + Args: + train_dataset: The training dataset to create the resample iterator from. + + Returns: + The resample data iterator (first element of the iterator tuple). + """ + args = self.args + resample_seed = getattr(args, 'seed', 42) + 1 + try: + set_random_seed( + resample_seed, + args.data_parallel_random_init, + args.te_rng_tracker, + ) + resample_data_iterator = self._prepare_data_iterator(train_dataset, use_origin_cyclic=True)[0] + finally: + set_random_seed( + args.seed, + args.data_parallel_random_init, + args.te_rng_tracker, + ) + return resample_data_iterator + + def resample_encode_failed_inputs(self, inputs: List[Dict], max_resample_rounds: int = 10) -> List[Dict]: + """Attempt to encode each input. If encoding fails, resample until we have enough valid samples. + + Args: + inputs: List of input data samples + max_resample_rounds: Maximum number of resample rounds + + Returns: + List of successfully encoded input samples with the same length as inputs + """ + return resample_encode_failed_inputs( + self.template, + self.resample_data_iterator, + inputs, + max_resample_rounds=max_resample_rounds, + strip_response=False, + ) + + def _assemble_teacher_outputs(self, encoded_batches: List[Dict]) -> None: + for encoded_batch in encoded_batches: + parsed = encoded_batch.pop('_teacher_parsed') + teacher_model_inputs = encoded_batch['teacher_model_inputs'] + teacher_out = assemble_teacher_output( + parsed, + teacher_model_inputs=teacher_model_inputs, + topk=self.gkd_logits_topk, + template_padding_free=self.template.padding_free, + device=self.device, + ) + if teacher_out.labels is not None: + teacher_out.labels = torch.roll(teacher_out.labels, shifts=-1, dims=-1) + encoded_batch['teacher_output'] = teacher_out + + def _compute_teacher_logits(self, encoded_batches: List[Dict], vp_stage: Optional[int] = None) -> None: + if self.use_teacher_api: + self._assemble_teacher_outputs(encoded_batches) + return + if self._is_self_distillation: + # Self-distillation teacher == current student weights. Computing it here (at batch + # preparation, once per steps_per_generation cycle) would reuse stale weights across the + # cycle's train steps. Defer to _replace_data_iterator so each train step recomputes the + # teacher with up-to-date student weights (weights are constant within a train step). + return + self._compute_teacher_logits_local(encoded_batches, vp_stage) + + def _compute_teacher_logits_local(self, encoded_batches: List[Dict], vp_stage: Optional[int] = None) -> None: + """Compute teacher_output for each micro-batch via a local forward. + + Handles both a separate fixed teacher and self-distillation (teacher == current student + weights). For self-distillation the caller is responsible for invoking this per train step + so the weights are current. + """ + topk = self.gkd_logits_topk + if self._is_self_distillation: + teacher_model = self.unwrapped_models[vp_stage or 0] + adapter_contexts = [] + if self._teacher_use_disable_adapter: + adapter_contexts = [m.disable_adapter() for m in self.peft_models] + outer_context = ContextManagers(adapter_contexts) + else: + teacher_model = self.teacher_models[vp_stage or 0] + outer_context = self.load_teacher_model_context() + + with torch.no_grad(), outer_context: + for encoded_batch in encoded_batches: + teacher_model_inputs = encoded_batch['teacher_model_inputs'] + teacher_batch = { + k: v.clone() if isinstance(v, torch.Tensor) else v + for k, v in teacher_model_inputs.items() + } + teacher_data = self._prepare_batch(teacher_batch, vp_stage) + teacher_data.pop('loss_scale', None) + teacher_labels = teacher_data.pop('labels', None) + teacher_logits = forward_step_helper(teacher_model, teacher_data) + if teacher_logits is not None: + teacher_logits = teacher_logits.detach() + + if topk is not None and teacher_logits is not None: + topk_logits, topk_indices = vocab_parallel_topk(teacher_logits, k=topk) + teacher_out = TeacherOutput(topk_logprobs=topk_logits, topk_indices=topk_indices) + else: + teacher_out = TeacherOutput(full_logits=teacher_logits) + + teacher_out.labels = teacher_labels + encoded_batch['teacher_output'] = teacher_out + + def _generate_and_score_completions(self, inputs: List[Dict]) -> List[Dict]: + """Unified rollout → teacher → encode pipeline (mirrors Megatron GRPO). + + Stages: determine data source → to_samples → (student) generate → teacher + requests/logprobs → encode micro-batches → teacher logits. Returns the flat + list of encoded micro-batches (length == total microbatches). + """ + data_source = self._determine_data_source() + + # Convert to samples (resample operates on dict, to_samples after) + samples = self.to_samples(inputs) + + if data_source == DataSource.STUDENT: + local_batch = self._get_local_rollout_batch(samples) + local_batch = self._generate_completions(local_batch) + samples = self._gather_rollout_results(local_batch) + self._log_completions_from_samples(samples) + + # Teacher API: build requests from samples, fetch logprobs + local_parsed = None + if self.use_teacher_api: + build_opsd_samples(samples) + teacher_requests = self._build_teacher_requests(samples) + if teacher_requests: + local_parsed = fetch_teacher_parsed_by_routing( + samples, + teacher_requests, + self.teacher_configs, + self.teacher_clients, + gather_fn=self._gather_teacher_requests, + infer_fn=lambda handle, client: self._infer_teacher_requests( + handle, topk=self.gkd_logits_topk, teacher_client=client), + scatter_fn=self._scatter_teacher_parsed, + is_main_process=self.is_main_process, + tag_key=self.args.teacher_tag_key) + + micro_batch_size = self.args.micro_batch_size + total_microbatches = len(samples) // micro_batch_size + assert total_microbatches * micro_batch_size == len(samples) + all_encoded_batches = [] + for i in range(total_microbatches): + start_idx = i * micro_batch_size + end_idx = start_idx + micro_batch_size + sample_slice = samples[start_idx:end_idx] + encoded_batch = self._encode_samples(sample_slice) + encoded_batch['data_source'] = data_source + if local_parsed is not None: + encoded_batch['_teacher_parsed'] = local_parsed[start_idx:end_idx] + all_encoded_batches.append(encoded_batch) + self._compute_teacher_logits(all_encoded_batches) + return all_encoded_batches + + def _build_rollout_buffer(self, data_iterator): + num_microbatches = self.args.num_microbatches + num_gen_steps = self.steps_per_generation if self.unwrapped_models[0].training else 1 + total_microbatches = num_microbatches * num_gen_steps + global_batch = [] + for _ in range(total_microbatches): + raw_batch = next(data_iterator) + if self.truncation_strategy == 'delete' and self.resample_data_iterator is not None: + raw_batch = self.resample_encode_failed_inputs(raw_batch) + global_batch.extend(raw_batch) + + all_encoded_batches = self._generate_and_score_completions(global_batch) + return [all_encoded_batches[i * num_microbatches:(i + 1) * num_microbatches] for i in range(num_gen_steps)] + + def _on_train_step_batch(self, encoded_batches): + # Self-distillation teacher == current student weights. Recompute per train step (weights are + # constant within a step) instead of once per generation cycle, so it tracks student updates + # across steps_per_generation. Runs outside the pipeline schedule, so PP > 1 is supported. + if self._is_self_distillation: + self._compute_teacher_logits_local(encoded_batches) + + def loss_func(self, + output_tensor: torch.Tensor, + *, + labels: torch.Tensor, + teacher_output: TeacherOutput, + data_source: DataSource = DataSource.DATASET): + """Compute GKD loss (JSD + optional SFT loss).""" + student_logits = output_tensor + + jsd_total, jsd_num_valid = gkd_loss( + student_logits, + teacher_output, + labels, + self.beta, + self.temperature, + gather_fn=tp_gather_topk, + log_softmax_fn=vocab_parallel_log_softmax, + kl_div_fn=vocab_parallel_kl_div) + jsd_loss_val = cp_reduce(jsd_total, jsd_num_valid, cp_size=self.args.context_parallel_size) + + loss = jsd_loss_val + + # Add SFT loss if enabled (skip for student-generated responses) + sft_loss = None + if self.sft_alpha > 0 and data_source != DataSource.STUDENT: + args = self.args + logits_sbv = student_logits.transpose(0, 1).contiguous() + model = self.unwrapped_models[0] + if hasattr(model, 'language_model'): + model = model.language_model + per_token_loss = model.compute_language_model_loss(labels, logits_sbv) + loss_mask = labels != -100 + sft_loss_sum = (per_token_loss * loss_mask).sum() + sft_loss_count = loss_mask.sum().float() + + # All-reduce across CP group for correct averaging + if args.context_parallel_size > 1: + sft_stats = torch.stack([sft_loss_sum, sft_loss_count]) + torch.distributed.all_reduce( + sft_stats, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group()) + sft_loss_sum, sft_loss_count = sft_stats[0], sft_stats[1] + + sft_loss = sft_loss_sum / sft_loss_count + + loss = loss + self.sft_alpha * sft_loss + + metric = {'loss': loss.detach().clone()} + if sft_loss is not None: + metric['jsd_loss'] = jsd_loss_val.detach().clone() + metric['sft_loss'] = sft_loss.detach().clone() + metric = self._all_reduce_metric(metric) + + loss = loss / mpu.get_context_parallel_world_size() + + # Flush completion logs at generation cycle boundaries. + if (self._step - 1) % self.steps_per_generation == 0: + self._flush_log_completions() + + return loss, metric + + def forward_step(self, data_iterator, model): + unwrapped_model = model.module.module + input_tensor = unwrapped_model.get_input_tensor() + vp_stage = unwrapped_model.vp_stage + + data = next(data_iterator) + data_source = data.pop('data_source', DataSource.DATASET) + teacher_output = data.pop('teacher_output') + data.pop('teacher_model_inputs', None) # consumed by _compute_teacher_logits; not needed for student forward + data = self._prepare_batch(data, vp_stage) + + data.pop('loss_scale', None) + labels = data.pop('labels', None) + + if input_tensor is not None: + unwrapped_model.set_input_tensor(input_tensor) + student_output = model(**data) + + return student_output, partial( + self.loss_func, + labels=labels, + teacher_output=teacher_output, + data_source=data_source, + ) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/gkd_utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/gkd_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f778eb6bc35910e80c277b7a15bfd8bdbdd08ceb --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/gkd_utils.py @@ -0,0 +1,67 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Megatron-specific GKD utilities: TP-aware gather/topk and CP reduce.""" +import torch +from megatron.core import mpu + + +def vocab_parallel_topk(logits: torch.Tensor, k: int) -> tuple: + """Global top-k from vocab-parallel sharded logits. TP=1 → plain torch.topk.""" + tp_size = mpu.get_tensor_model_parallel_world_size() + if tp_size == 1: + return torch.topk(logits, k=k, dim=-1) + + tp_rank = mpu.get_tensor_model_parallel_rank() + tp_group = mpu.get_tensor_model_parallel_group() + partition_vocab_size = logits.shape[-1] + + local_topk_vals, local_topk_ids = torch.topk(logits, k=k, dim=-1) + local_topk_ids = local_topk_ids + tp_rank * partition_vocab_size + + gathered_vals = [torch.empty_like(local_topk_vals) for _ in range(tp_size)] + gathered_ids = [torch.empty_like(local_topk_ids) for _ in range(tp_size)] + torch.distributed.all_gather(gathered_vals, local_topk_vals, group=tp_group) + torch.distributed.all_gather(gathered_ids, local_topk_ids, group=tp_group) + + all_vals = torch.cat(gathered_vals, dim=-1) + all_ids = torch.cat(gathered_ids, dim=-1) + global_topk_vals, sel = torch.topk(all_vals, k=k, dim=-1) + global_topk_ids = torch.gather(all_ids, dim=-1, index=sel) + return global_topk_vals, global_topk_ids + + +def tp_gather_topk(logits: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: + """Gather logits at global top-k indices with TP-aware partitioning.""" + tp_size = mpu.get_tensor_model_parallel_world_size() + if tp_size == 1: + return torch.gather(logits, dim=-1, index=indices) + + tp_rank = mpu.get_tensor_model_parallel_rank() + partition_vocab_size = logits.shape[-1] + vocab_start = tp_rank * partition_vocab_size + + in_range = (indices >= vocab_start) & (indices < vocab_start + partition_vocab_size) + local_indices = (indices - vocab_start).clamp(0, partition_vocab_size - 1) + gathered = torch.gather(logits, dim=-1, index=local_indices) + gathered = gathered.masked_fill(~in_range, float('-inf')) + + gathered_for_reduce = gathered.detach() + torch.distributed.all_reduce( + gathered_for_reduce, op=torch.distributed.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group()) + return torch.where(in_range, gathered, gathered_for_reduce) + + +def cp_reduce(total_loss, num_valid, *, cp_size): + """Normalize total_loss by num_valid with CP all-reduce when cp_size > 1.""" + num_valid_f = num_valid.float() if isinstance(num_valid, torch.Tensor) else torch.tensor( + float(num_valid), device=total_loss.device) + if cp_size > 1: + torch.distributed.all_reduce( + num_valid_f, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group()) + torch.distributed.all_reduce( + total_loss, op=torch.distributed.ReduceOp.SUM, group=mpu.get_context_parallel_group()) + # Avoid the host-device sync from ``if num_valid_f == 0`` (num_valid_f is a GPU scalar, + # so a Python bool() forces a .item() stream sync every step). Keep the zero-valid + # guard on-device with torch.where: num_valid==0 -> loss 0, else total_loss/num_valid. + safe = torch.where(num_valid_f == 0, torch.ones_like(num_valid_f), num_valid_f) + loss = total_loss / safe + return torch.where(num_valid_f == 0, torch.zeros_like(loss), loss) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/grpo_trainer.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/grpo_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..92b4ea2163d108b4334876e1b67e955156e0ddc4 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/grpo_trainer.py @@ -0,0 +1,1632 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from accelerate.utils import broadcast_object_list +from collections import defaultdict, deque +from contextlib import contextmanager +from copy import copy, deepcopy +from functools import partial +from mcore_bridge import set_random_seed +from megatron.core import mpu +from typing import Any, Dict, List, Optional, Tuple, Union + +from swift.infer_engine.protocol import RequestConfig, RolloutInferRequest, RolloutOutput +from swift.megatron.arguments import MegatronArguments +from swift.megatron.utils import RouterReplayHelper, get_padding_to, set_router_replay_data +from swift.rl_core.advantage import (compute_advantages, compute_reward_metrics, compute_teacher_kl_per_token, + expand_advantage_to_per_token) +from swift.rl_core.data import GRPOBatch, GRPOSample +from swift.rl_core.grpo_algorithm import score_completions +from swift.rl_core.resample import resample_encode_failed_inputs +from swift.rlhf_trainers.gkd_helpers import (assemble_teacher_completion_logprobs, build_opsd_samples, + build_teacher_requests, encode_teacher_view, + fetch_teacher_parsed_by_routing, remap_teacher_logps_to_student_frame, + should_compute_local_teacher_logps) +from swift.rlhf_trainers.grpo_trainer import DataType +from swift.rlhf_trainers.utils import (collate_to_grpo_micro_batch, encode_sample, make_reward_weights, + pad_logps_back_to_batch, profiling_context, profiling_decorator, + resolve_reward_funcs) +from swift.template import Template +from swift.utils import get_logger +from .rlhf_mixin import MegatronRLHFTrainer +from .rollout_mixin import MegatronRolloutMixin +from .utils import gather, gather_object, reconstruct_tensor_cp +from .vocab_parallel_utils import compute_logps_and_entropy_from_logits + +try: + from megatron.core.transformer.moe.router_replay import RouterReplay, RouterReplayAction +except ImportError: + RouterReplay = None + RouterReplayAction = None + +logger = get_logger() + + +class MegatronGRPOTrainer(MegatronRolloutMixin, MegatronRLHFTrainer): + + # Per-sample container class used by MegatronRolloutMixin.to_samples. + sample_cls = GRPOSample + + def __init__(self, args: MegatronArguments, template: Template, **kwargs): + self.vllm_client = kwargs.pop('vllm_client') + self.args = args + self._setup_teacher() + super().__init__(args, template) + self.args = args + self.hf_model_dir = args.model_info.model_dir + self.processing_class = self.template.processor + self._prepare_metrics() + self._init_grpo_params() + self._init_rollout_engine() + self._prepare_rewards() + self.resample_data_iterator = None + + def prepare_model(self): + super().prepare_model() + self._load_teacher_model() + + def train(self, train_dataset, val_dataset): + if self.dynamic_sample or self.truncation_strategy == 'delete': + self.resample_data_iterator = self._init_resample_data_iterator(train_dataset) + super().train(train_dataset, val_dataset) + + def _init_grpo_params(self): + """Initialize GRPO-specific parameters. + + Note: Common rollout params (world_size, process_index, device, request_config, etc.) + are initialized by MegatronRolloutMixin._init_rollout_params(). + """ + args: MegatronArguments = self.args + + # GRPO algorithm params + self.num_generations = args.num_generations # G in the GRPO paper + self.num_generations_eval = args.num_generations_eval or self.num_generations + self.beta = args.beta + self.loss_type = args.loss_type + self.epsilon_low = args.epsilon + self.epsilon_high = args.epsilon_high if args.epsilon_high is not None else args.epsilon + self.top_entropy_quantile = args.top_entropy_quantile + self.importance_sampling_level = args.importance_sampling_level + + # SAPO, https://arxiv.org/abs/2511.20347 + self.tau_pos = args.tau_pos + self.tau_neg = args.tau_neg + + # REAL, https://arxiv.org/abs/2602.05630 + self.real_tau = args.real_tau + + # FIPO, https://arxiv.org/abs/2603.19835 + self.fipo_gamma = 2**(-1 / args.fipo_decay_rate) + self.fipo_clip_range = args.fipo_clip_range + self.fipo_clip_high_only = args.fipo_clip_high_only + self.fipo_safety_threshold = args.fipo_safety_threshold + + # DAPO, https://arxiv.org/abs/2503.14476 + self.dynamic_sample = args.dynamic_sample + self.max_resample_times = args.max_resample_times + self.overlong_filter = args.overlong_filter + + # Dr. GRPO / RLOO / REINFORCE++ + self.scale_rewards = args.scale_rewards + self.advantage_estimator = args.advantage_estimator + self.kl_in_reward = args.kl_in_reward + if self.scale_rewards == 'gdpo' and self.kl_in_reward: + logger.warning('GDPO mode does not support kl_in_reward=True. Setting kl_in_reward=False.') + self.kl_in_reward = False + + self.log_entropy = args.log_entropy + self.compute_entropy = self.log_entropy or self.top_entropy_quantile < 1.0 + + # Rollout Importance Sampling Correction + self.rollout_importance_sampling_mode = args.rollout_importance_sampling_mode + self.rollout_importance_sampling_threshold = args.rollout_importance_sampling_threshold + self.log_rollout_offpolicy_metrics = args.log_rollout_offpolicy_metrics + + # Off-Policy Sequence Masking + self.off_policy_sequence_mask_delta = args.off_policy_sequence_mask_delta + + # batch size (completion-level) + self.generation_batch_size = args.generation_batch_size + self.steps_per_generation = args.steps_per_generation + self.global_batch_size = args.global_batch_size + self.micro_batch_size = args.micro_batch_size + self.per_device_generation_batch_size = args.per_device_generation_batch_size + + # truncation_strategy support + self.truncation_strategy = args.truncation_strategy + + self.teacher_kl_coef = args.teacher_kl_coef + + def _init_rollout_engine(self): + """Initialize rollout engine with GRPO-specific extensions.""" + super()._init_rollout_engine() + + # GRPO-specific initialization + self.async_generate = self.args.async_generate # TODO + self.use_gym_env = self._resolve_use_gym_env() + self._buffered_inputs = None + + # Rollout importance sampling requires vLLM >= 0.10.2 + self.disable_rollout_importance_sampling = not self.vllm_version_ge_0_10_2 + if not self.vllm_version_ge_0_10_2 and getattr(self.args, 'rollout_importance_sampling_mode', None) is not None: + raise ValueError('rollout_importance_sampling_mode is not supported in vLLM version < 0.10.2, ' + 'please update vLLM to 0.10.2 or later.') + + def _resolve_use_gym_env(self) -> bool: + """Resolve `use_gym_env` for the trainer, mirroring the HF rollout_mixin logic. + + Priority: + 1. Explicit `args.use_gym_env` (works for both server and colocate). + 2. In server mode, auto-detect from the connected vLLM rollout server. + 3. Default to False. + """ + args = self.args + explicit = getattr(args, 'use_gym_env', None) + if explicit is not None: + return bool(explicit) + + if not getattr(self, 'use_vllm', False) or self.vllm_mode != 'server': + return False + + # super()._init_rollout_engine() has already called vllm_client.get_engine_type() + # on the main rank, so vllm_client.use_gym_env is populated there. + if self.is_main_process: + value = [bool(getattr(self.vllm_client, 'use_gym_env', False))] + else: + value = [False] + return broadcast_object_list(value, from_process=self.world_size - 1)[0] + + def _prepare_rewards(self): + args = self.args + reward_funcs_cfg = args.reward_funcs.copy() + if not isinstance(reward_funcs_cfg, list): + reward_funcs_cfg = [reward_funcs_cfg] + + self.reward_funcs, self.reward_func_names = resolve_reward_funcs(reward_funcs_cfg, args=self.args) + + # use_gym_env: gym total_reward is appended as an extra reward column so it can + # blend with reward_funcs via reward_weights. When reward_funcs is empty, it becomes + # the single reward source. + if self.use_gym_env: + self.reward_func_names.append('gym_reward') + + self.reward_weights = make_reward_weights(args.reward_weights, len(self.reward_func_names), self.device) + + self.reward_model_plugins = [None] * len(self.reward_funcs) + + assert self.reward_funcs or self.use_gym_env or self._has_teacher, \ + 'reward_funcs is not set (or pass --use_gym_env true / a teacher for OPD-RL)' + + def _init_resample_data_iterator(self, train_dataset): + """Initialize an independent data iterator for resampling. + + Uses a different seed (args.seed + 1) to avoid overlapping with training samples. + + Args: + train_dataset: The training dataset to create the resample iterator from. + + Returns: + The resample data iterator (first element of the iterator tuple). + """ + args = self.args + resample_seed = getattr(args, 'seed', 42) + 1 + try: + set_random_seed( + resample_seed, + args.data_parallel_random_init, + args.te_rng_tracker, + ) + # TODO: VPP (Virtual Pipeline Parallelism) + resample_data_iterator = self._prepare_data_iterator(train_dataset, use_origin_cyclic=True)[0] + finally: + set_random_seed( + args.seed, + args.data_parallel_random_init, + args.te_rng_tracker, + ) + return resample_data_iterator + + def _build_rollout_buffer(self, data_iterator): + num_gen_steps = self.steps_per_generation if self.unwrapped_models[0].training else 1 + num_iters_per_step = self.get_num_iters_per_step() + rollout_batch = [] + for _ in range(num_iters_per_step): + rollout_batch.extend(next(data_iterator)) + micro_batch_data = self._generate_and_score_completions(rollout_batch) + num_mini_batch = self.global_batch_size // (self.micro_batch_size * mpu.get_data_parallel_world_size()) + mini_batch_data = [ + micro_batch_data[i:i + num_mini_batch] for i in range(0, len(micro_batch_data), num_mini_batch) + ] + assert len(mini_batch_data) == num_gen_steps + return mini_batch_data + + def _generate_and_score_completions(self, inputs: DataType): + # Get or create the rollout group (TP×PP×CP) + + rollout_group = self._get_rollout_group() + + # Resample for encoding failed data when truncation_strategy is 'delete' + # This handles: (1) prompt length exceeds max_length, (2) multimodal encoding failures + # Do this before get_local_rollout_batch to process prompt-level data. + # Resample operates on dict rows; build per-sample objects only once the + # training data is finalized (mirrors HF / Ray GRPO). + if self.truncation_strategy == 'delete': + inputs = self.resample_encode_failed_inputs(inputs) + samples = self.to_samples(inputs) + + samples = self.get_local_rollout_batch(samples) + + samples = self._generate_completions(samples) + + rewards_per_func = self._score_completions(samples) + + # Dynamic sampling for std=0 groups (DAPO) + if self.dynamic_sample: + samples, rewards_per_func = self._dynamic_sampling(samples, rewards_per_func) + + # Log completions after all filtering so they align with rewards/advantages (issue #9533). + self._log_completions_from_samples(samples) + + # Gather rollout data across rollout group + total_samples = gather_object(samples, group=rollout_group) + mini_batch_data = [] + template = self.template + + # OPSD: build each sample's teacher view when this batch uses privileged teacher_prompt. + has_opsd_batch = build_opsd_samples(total_samples) + is_opsd = (not self.use_teacher_api and has_opsd_batch + and (self._has_teacher_explicit or self._is_dynamic_self_distillation)) + + # Step 1: Encode batches and compute logps first (unified flow like GRPOTrainer) + with self._template_context(template): + for s in total_samples: + s.encoded = encode_sample(s, template) + s.encoded.pop('_extra_kwargs', None) + + for idx in range(0, len(total_samples), self.micro_batch_size): + sample_slice = total_samples[idx:idx + self.micro_batch_size] + model_inputs, grpo_batch = collate_to_grpo_micro_batch( + sample_slice, + template, + device=self.device, + padding_to=get_padding_to(self.args), + router_replay_mode=self.args.router_replay_mode, + ) + teacher_inputs = None + if is_opsd: + teacher_inputs = self._collate_teacher_opsd_batch(sample_slice, template) + # Wire format: model_inputs (dict) + grpo_batch (consumed by forward_step) + data = model_inputs + data['grpo_batch'] = grpo_batch + with profiling_context(self, 'compute_ref_old_logps'): + data = self._maybe_compute_logps(data, teacher_inputs=teacher_inputs) + mini_batch_data.append(data) + + if self._has_teacher and self.use_teacher_api: + build_opsd_samples(total_samples) # OPSD: populate teacher_messages for build_teacher_requests + self._assemble_teacher_api_logps(total_samples, mini_batch_data) + + # Step 2: Compute KL from logps if kl_in_reward is enabled + kl_values = None + if self.kl_in_reward and self.beta != 0.0: + kl_values = self._compute_kl_from_batches(mini_batch_data) + + # Step 3: Compute the per-sequence base advantage (with ref-KL penalty if kl_in_reward). + advantages = self._compute_advantages(samples, rewards_per_func, kl_values=kl_values) + total_advantages = gather(advantages, group=rollout_group) + + # Step 4: Write the advantage onto each batch, expanding the per-sequence base advantage to + # per-token [B, T] here so the OPD-RL signed teacher log-ratio is added per token + # (adv_t = base + coef * (teacher_logp - student_logp)). + for idx, micro_batch_encoded in enumerate(mini_batch_data): + grpo_batch = micro_batch_encoded['grpo_batch'] + start_idx = idx * self.micro_batch_size + end_idx = start_idx + grpo_batch.completion_mask.shape[0] + micro_batch_advantages = total_advantages[start_idx:end_idx] + grpo_batch.advantages = expand_advantage_to_per_token( + micro_batch_advantages, + grpo_batch.completion_mask, + teacher_per_token_logps=grpo_batch.teacher_per_token_logps, + policy_per_token_logps=grpo_batch.old_per_token_logps + if grpo_batch.teacher_per_token_logps is not None else None, + teacher_kl_coef=self.teacher_kl_coef if grpo_batch.teacher_per_token_logps is not None else 0.0, + ) + if any(m['grpo_batch'].teacher_per_token_logps is not None for m in mini_batch_data): + self._log_teacher_kl_metric(mini_batch_data) + + if self.loss_type in ['cispo', 'dapo', 'fipo']: + # Calculate num_items_in_batch + # Count completion tokens from all mini_batch_data (this includes gathered data from rollout_group) + # Use completion_mask.sum() for both padding_free and non-padding_free modes + # since we want the count of actual completion tokens, not sequence lengths + total_token_count = sum(batch_data['grpo_batch'].completion_mask.sum().item() + for batch_data in mini_batch_data) + + # All-reduce across all ranks + total_token_count_tensor = torch.tensor(total_token_count, dtype=torch.int, device=self.device) + torch.distributed.all_reduce(total_token_count_tensor) + + # Divide by rollout_group_size to account for duplicate counting within each rollout_group + # Each rollout_group (TP×PP×CP ranks) has the same gathered data, so we need to normalize + rollout_group_size = ( + mpu.get_tensor_model_parallel_world_size() * mpu.get_pipeline_model_parallel_world_size() + * mpu.get_context_parallel_world_size()) + num_items_in_batch = total_token_count_tensor / rollout_group_size + # Store num_items_in_batch in each mini_batch_data for token-normalized losses + for batch_data in mini_batch_data: + batch_data['grpo_batch'].num_items_in_batch = num_items_in_batch + + return mini_batch_data + + def _compute_fipo_influence(self, log_ratio: torch.Tensor, coef_1: torch.Tensor, advantages: torch.Tensor, + completion_mask: torch.Tensor) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]: + """Compute FIPO token-level influence weight from discounted Future-KL.""" + future_kl_delta = log_ratio.masked_fill(~completion_mask, 0.0) + + if self.args.delta is not None: + delta = torch.as_tensor(self.args.delta, dtype=log_ratio.dtype, device=log_ratio.device) + high_ratio_mask = coef_1 > delta + future_kl_delta = torch.where(high_ratio_mask, torch.zeros_like(future_kl_delta), future_kl_delta) + + seq_len = future_kl_delta.shape[1] + future_kl = torch.zeros_like(future_kl_delta) + positions = torch.arange(seq_len, device=log_ratio.device).unsqueeze(1) + gamma = torch.as_tensor(self.fipo_gamma, dtype=log_ratio.dtype, device=log_ratio.device) + chunk_size = 128 + for chunk_start in range(0, seq_len, chunk_size): + chunk_end = min(seq_len, chunk_start + chunk_size) + chunk_positions = torch.arange(chunk_start, chunk_end, device=log_ratio.device).unsqueeze(0) + distance = chunk_positions - positions + future_mask = distance >= 0 + decay_block = torch.pow(gamma, distance.clamp(min=0)) * future_mask.to(log_ratio.dtype) + future_kl += torch.matmul(future_kl_delta[:, chunk_start:chunk_end], decay_block.t()) + future_kl = future_kl.masked_fill(~completion_mask, 0.0) + + influence_weight = torch.exp(future_kl) + + if self.fipo_clip_range: + high = 1 + self.fipo_clip_range + low = 1.0 if self.fipo_clip_high_only else 1 - self.fipo_clip_range + influence_weight = torch.clamp(influence_weight, min=low, max=high) + influence_weight = influence_weight.detach() + + safety_mask = torch.ones_like(completion_mask, dtype=torch.bool) + if self.fipo_safety_threshold is not None: + negative_advantage = advantages < 0 + high_is_ratio = coef_1 > self.fipo_safety_threshold + safety_mask = ~(negative_advantage & high_is_ratio) + influence_weight = torch.where(safety_mask, influence_weight, + torch.clamp(influence_weight, min=0.8, max=1.0)) + + metrics = { + 'future_kl': future_kl, + 'influence_weight': influence_weight, + 'safety_mask': safety_mask, + } + return influence_weight, metrics + + def _build_log_table(self) -> Dict[str, list]: + """Extend base table with GRPO-specific rewards and advantages.""" + table = super()._build_log_table() + table.update({k: list(v) for k, v in self._logs['rewards'].items()}) + table['advantages'] = list(self._logs['advantages']) + return table + + def _get_request_config(self) -> RequestConfig: + request_config = copy(self.request_config) + if self.args.vllm_mode == 'colocate' and self.vllm_tensor_parallel_size > 1: + # Set request_config.seed + # 1. Ensure that the seed for vLLM Engines within each TP (Tensor Parallelism) group is the same; + # otherwise, the program may hang. + # 2. Ensure that the seed for vLLM Engines across different TP groups is different; + # otherwise, identical completions will be generated. + batch_size = self.per_device_generation_batch_size + batch_size *= self.vllm_tensor_parallel_size + # Since the TP (Tensor Parallelism) group gathers the inputs, + # multiply the batch size by the TP parallel size. + request_config.seed = batch_size * (self.process_index // self.vllm_tensor_parallel_size) + + return request_config + + def _server_rollout(self, + inputs: DataType, + request_config: RequestConfig, + is_global_inputs: bool = False) -> List[RolloutOutput]: + # TODO: async generate + infer_requests = self.samples2requests(inputs) + + if is_global_inputs: + per_device_size = len(infer_requests) // self.world_size + all_requests = infer_requests + all_requests_lengths = [per_device_size] + [0] * (self.world_size - 1) + else: + all_requests = gather_object(infer_requests) + all_requests_lengths = gather_object([len(infer_requests)]) + + if not any(requests for requests in all_requests): + return [] + + if self.is_main_process: + all_outputs: List[RolloutOutput] = self.vllm_client.infer( + infer_requests=all_requests, request_config=request_config) + assert len(all_outputs) == len(all_requests) # TODO: dynamic num of samples + else: + all_outputs = [None] * len(all_requests) + + if not is_global_inputs: + all_outputs = broadcast_object_list(all_outputs, from_process=self.world_size - 1) + start_idx = sum(all_requests_lengths[:self.process_index]) + end_idx = start_idx + all_requests_lengths[self.process_index] + outputs = all_outputs[start_idx:end_idx] + else: + outputs = all_outputs if self.is_main_process else [] + return outputs + + @profiling_decorator + def _score_completions(self, samples: List[GRPOSample]) -> torch.Tensor: + """Score completions using all reward functions. + + Args: + samples: List of on-policy samples carrying messages + rollout fields. + + Returns: + rewards_per_func: Tensor of shape (num_examples, num_reward_funcs) with local reward values + """ + return score_completions( + samples, + reward_funcs=self.reward_funcs, + reward_model_plugins=self.reward_model_plugins, + use_gym_env=self.use_gym_env, + device=self.device, + trainer_state=self.state, + ) + + def _compute_advantages(self, + samples: List[GRPOSample], + rewards_per_func: torch.Tensor, + kl_values: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Compute the per-sequence base advantage for RL training. + + The OPD-RL teacher signal is not injected here; it is applied per-token when the base + advantage is expanded to ``[B, T]`` onto each batch (see ``expand_advantage_to_per_token``). + + Args: + samples: Local on-policy samples (only the count is used for slicing) + rewards_per_func: Reward per function for local data samples + kl_values: Optional KL values for kl_in_reward mode, shape [total_samples] + + Returns: + advantages: Computed advantages for local batch, shape [local_batch_size] + """ + mode = 'train' if self.unwrapped_models[0].training else 'eval' + assert len(samples) == rewards_per_func.shape[0] + num_generations = self.num_generations if mode == 'train' else self.num_generations_eval + + # Gather local rewards into a global tensor: GRPO group normalization needs every + # completion of the same prompt visible on one rank (groups span DP ranks). + # ``kl_values`` is already global (gathered in ``_compute_kl_from_batches``). + # OPD-RL pure distillation: no reward_funcs -> a [N, 0] tensor. + if rewards_per_func.shape[1] == 0: + global_count = sum(gather_object([rewards_per_func.shape[0]])) + total_rewards_per_func = torch.zeros((global_count, 0), dtype=torch.float32, device=self.device) + else: + total_rewards_per_func = gather(rewards_per_func) + + advantages, weighted_rewards = compute_advantages( + rewards_per_func=total_rewards_per_func, + reward_weights=self.reward_weights, + num_generations=num_generations, + advantage_estimator=self.advantage_estimator, + scale_rewards=self.scale_rewards, + kl_in_reward=self.kl_in_reward, + beta=self.beta, + kl_values=kl_values, + ) + + reward_metrics = compute_reward_metrics( + rewards=weighted_rewards, + rewards_per_func=total_rewards_per_func, + reward_func_names=self.reward_func_names, + num_generations=num_generations, + scale_rewards=self.scale_rewards, + ) + self._metrics[mode]['reward'].append(reward_metrics.reward_mean) + self._metrics[mode]['reward_std'].append(reward_metrics.reward_std) + self._metrics[mode]['frac_reward_zero_std'].append(reward_metrics.frac_reward_zero_std) + if kl_values is not None: + self._metrics[mode]['kl'].append(kl_values.nanmean().item()) + for name in self.reward_func_names: + self._metrics[mode][f'rewards/{name}/mean'].append(reward_metrics.per_func_mean[name]) + self._metrics[mode][f'rewards/{name}/std'].append(reward_metrics.per_func_std[name]) + self._logs['advantages'].extend(advantages.tolist()) + for i, name in enumerate(self.reward_func_names): + self._logs['rewards'][name].extend(total_rewards_per_func[:, i].tolist()) + + # Slice the global advantages back to this rank's local samples. + slice_start = self.process_index * len(samples) + slice_end = slice_start + len(samples) + return advantages[slice_start:slice_end] + + def _dynamic_sampling(self, rollout_batch: List[GRPOSample], + rewards_per_func: torch.Tensor) -> Tuple[List[GRPOSample], torch.Tensor]: + """ + Perform dynamic sampling to replace samples with zero-reward-variance groups. + + This method implements DAPO (https://arxiv.org/abs/2503.14476) by replacing + samples from groups with zero reward variance (std=0) through resampling. + + Args: + rollout_batch: local rollout data samples + rewards_per_func: reward per function for local data samples + rollout_group: rollout communication group + + Returns: + tuple: (rollout_batch, rewards_per_func) with zero-variance groups replaced by resampled data + """ + resample_count = 0 + valid_samples = [] + valid_rewards_per_func = [] + origin_data = (rollout_batch, rewards_per_func) + + while resample_count < self.max_resample_times: + # Gather all samples and rewards across rollout group first + global_rollout_batch = gather_object(rollout_batch) + if rewards_per_func.shape[1] == 0: + global_rewards_per_func = torch.zeros((len(global_rollout_batch), 0), + dtype=torch.float32, + device=self.device) + else: + global_rewards_per_func = gather(rewards_per_func) + + # Compute reward std for the entire global batch + # We need to compute std on the gathered data to get a global mask + global_rewards = (global_rewards_per_func * self.reward_weights.unsqueeze(0)).nansum(dim=1) + mode = 'train' if self.unwrapped_models[0].training else 'eval' + num_generations = self.num_generations if mode == 'train' else self.num_generations_eval + grouped_rewards = global_rewards.view(-1, num_generations) + # Handle edge case when num_generations=1 + if num_generations > 1: + group_rewards_std = grouped_rewards.std(dim=1).repeat_interleave(num_generations) + else: + group_rewards_std = torch.zeros_like(global_rewards) + global_valid_mask = (group_rewards_std > 0) + + # Filter valid samples based on std > 0 + valid_samples.extend([sample for sample, mask in zip(global_rollout_batch, global_valid_mask) if mask]) + valid_rewards_per_func.append(global_rewards_per_func[global_valid_mask]) + + if len(valid_samples) >= self.generation_batch_size: + break + + num_iters_per_step = self.get_num_iters_per_step() + next_rollout_prompt_batch = [] + for _ in range(num_iters_per_step): + next_rollout_prompt_batch.extend(next(self.resample_data_iterator)) + + # Resample for encoding failed data when truncation_strategy is 'delete' + if self.truncation_strategy == 'delete': + next_rollout_prompt_batch = self.resample_encode_failed_inputs(next_rollout_prompt_batch) + + # Convert dict rows to per-sample objects before rollout + next_rollout_prompt_batch = self.to_samples(next_rollout_prompt_batch) + + # Repeat num_generations times and get local slice + rollout_batch = self.get_local_rollout_batch(next_rollout_prompt_batch) + + # Generate and score new completions + rollout_batch = self._generate_completions(rollout_batch) + rewards_per_func = self._score_completions(rollout_batch) + resample_count += 1 + + if len(valid_samples) >= self.generation_batch_size: + # Get local slice of valid samples + rank = self.process_index + per_device_batch_size = self.per_device_generation_batch_size + data_slice = slice(rank * per_device_batch_size, (rank + 1) * per_device_batch_size) + rollout_batch = valid_samples[:self.generation_batch_size][data_slice] + rewards_per_func = torch.cat(valid_rewards_per_func)[:self.generation_batch_size][data_slice] + else: + logger.warning(f'There are still std=0 groups present after {self.max_resample_times} retries.') + rollout_batch, rewards_per_func = origin_data + + return rollout_batch, rewards_per_func + + def _maybe_compute_logps(self, batch: Dict[str, Any], teacher_inputs: Optional[Tuple] = None) -> Dict[str, Any]: + grpo_batch: GRPOBatch = batch.pop('grpo_batch') + seq_lengths = grpo_batch.seq_lengths + batch_size = grpo_batch.completion_mask.shape[0] + max_seq_len = grpo_batch.completion_mask.shape[1] + + # batch is now clean model forward kwargs (template.encode guarantees this) + inputs = batch + if self.beta != 0.0: + with self.null_ref_context() as ref_models: + assert len(ref_models) == 1, 'GRPO currently does not support VPP.' + ref_model = ref_models[0] + ref_per_token_logps_packed, _ = self.compute_per_token_logps( + ref_model, iter([deepcopy(inputs)]), temperature=self.temperature) + if self.template.padding_free: + ref_per_token_logps, _ = pad_logps_back_to_batch( + logps_rmpad=ref_per_token_logps_packed, + logits_to_keep=max_seq_len, + batch_size=batch_size, + seq_lengths=seq_lengths) + else: + ref_per_token_logps = ref_per_token_logps_packed + grpo_batch.ref_per_token_logps = ref_per_token_logps + + if should_compute_local_teacher_logps( + has_teacher_explicit=self._has_teacher_explicit, + is_dynamic_self_distillation=self._is_dynamic_self_distillation, + use_teacher_api=self.use_teacher_api, + has_opsd_batch=teacher_inputs is not None, + ): + grpo_batch.teacher_per_token_logps = self._compute_teacher_logps( + inputs, grpo_batch, teacher_inputs=teacher_inputs) + + if self.enable_routing_replay: + if self.args.router_replay_mode == 'R2': + RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD) + if self.args.router_replay_mode == 'R3': + RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + + old_per_token_logps_packed, routing_topk_idx = self.compute_per_token_logps( + self.unwrapped_models[0], iter([deepcopy(inputs)]), temperature=self.temperature) + if self.template.padding_free: + old_per_token_logps, _ = pad_logps_back_to_batch( + logps_rmpad=old_per_token_logps_packed, + logits_to_keep=max_seq_len, + batch_size=batch_size, + seq_lengths=seq_lengths) + else: + old_per_token_logps = old_per_token_logps_packed + grpo_batch.old_per_token_logps = old_per_token_logps + + if self.enable_routing_replay: + batch['routed_experts'] = routing_topk_idx + RouterReplay.clear_global_indices() + RouterReplay.clear_global_router_replay_action() + + batch['grpo_batch'] = grpo_batch + return batch + + def _collate_teacher_opsd_batch(self, sample_slice: List[GRPOSample], template) -> Tuple[Dict[str, Any], GRPOBatch]: + """Encode + collate the OPSD teacher view (teacher_prompt + same response). + + Produces a teacher micro-batch whose completion frame may differ in length from + the student's; ``_compute_teacher_logps`` remaps the teacher logps back onto the + student frame. The student ``encoded`` is restored so downstream logps/loss keep + the student frame. ``build_teacher_view`` is already done by ``build_opsd_samples``. + """ + for s in sample_slice: + s.encoded = encode_teacher_view(s, template) + teacher_model_inputs, teacher_grpo_batch = collate_to_grpo_micro_batch( + sample_slice, + template, + device=self.device, + padding_to=get_padding_to(self.args), + router_replay_mode=self.args.router_replay_mode, + ) + teacher_model_inputs['grpo_batch'] = teacher_grpo_batch + for s in sample_slice: + s.encoded = encode_sample(s, template) + s.encoded.pop('_extra_kwargs', None) + return teacher_model_inputs, teacher_grpo_batch + + def _compute_teacher_logps(self, + inputs: Dict[str, Any], + grpo_batch: GRPOBatch, + teacher_inputs: Optional[Tuple] = None) -> torch.Tensor: + """OPD-RL: per-token teacher logp on the sampled tokens via a local teacher forward. + + Reuses ``compute_per_token_logps`` (same path as old/ref logps) so the teacher logp + frame matches the policy's (token-in-token-out). Same-model LoRA self-distillation runs + the student under ``disable_adapter``; otherwise the separate teacher model is used. + For OPSD, the teacher forwards its own ``teacher_inputs`` (teacher_prompt + same + response) and the result is remapped onto the student's completion frame. + """ + is_opsd = teacher_inputs is not None + if is_opsd: + teacher_model_inputs, teacher_grpo_batch = teacher_inputs + fwd_inputs = {k: v for k, v in teacher_model_inputs.items() if k != 'grpo_batch'} + fwd_batch = teacher_grpo_batch + else: + fwd_inputs = inputs + fwd_batch = grpo_batch + seq_lengths = fwd_batch.seq_lengths + batch_size = fwd_batch.completion_mask.shape[0] + max_seq_len = fwd_batch.completion_mask.shape[1] + + if self._teacher_use_disable_adapter: + from contextlib import ExitStack + with ExitStack() as stack: + for m in self.peft_models: + stack.enter_context(m.disable_adapter()) + teacher_logps_packed, _ = self.compute_per_token_logps( + self.unwrapped_models[0], iter([deepcopy(fwd_inputs)]), temperature=self.temperature) + else: + # Dynamic self-distillation (teacher_models is None): teacher = student (same + # weights including LoRA). + models = self.teacher_models if self.teacher_models else self.unwrapped_models + with self.load_teacher_model_context(): + teacher_logps_packed, _ = self.compute_per_token_logps( + models[0], iter([deepcopy(fwd_inputs)]), temperature=self.temperature) + + if self.template.padding_free: + teacher_per_token_logps, _ = pad_logps_back_to_batch( + logps_rmpad=teacher_logps_packed, + logits_to_keep=max_seq_len, + batch_size=batch_size, + seq_lengths=seq_lengths) + else: + teacher_per_token_logps = teacher_logps_packed + if is_opsd: + teacher_per_token_logps = remap_teacher_logps_to_student_frame(teacher_per_token_logps, + teacher_grpo_batch.completion_mask, + grpo_batch.completion_mask) + return teacher_per_token_logps + + def _assemble_teacher_api_logps(self, total_samples: List[GRPOSample], mini_batch_data: List[Dict[str, + Any]]) -> None: + """OPD-RL teacher API: fetch the sampled token's logp per response position + (``prompt_logprobs=0``) and write it as ``teacher_per_token_logps`` (completion frame) + on each micro-batch's GRPOBatch. Same sampled-token semantics as HF/Ray OPD-RL. + + Each sample routes to exactly one teacher by tag (single teacher = all samples); + results are fetched globally then sliced per micro-batch. + """ + requests = build_teacher_requests(total_samples, self.template) + all_rti = [s.response_token_ids for s in total_samples] + parsed = fetch_teacher_parsed_by_routing( + total_samples, + requests, + self.teacher_configs, + self.teacher_clients, + gather_fn=self._gather_teacher_requests, + infer_fn=lambda handle, client: self._infer_teacher_requests(handle, topk=0, teacher_client=client), + scatter_fn=self._scatter_teacher_parsed, + is_main_process=self.is_main_process, + tag_key=self.args.teacher_tag_key) + + offset = 0 + for data in mini_batch_data: + grpo_batch: GRPOBatch = data['grpo_batch'] + device = grpo_batch.completion_mask.device + n = grpo_batch.completion_mask.shape[0] + teacher_out = assemble_teacher_completion_logprobs( + parsed[offset:offset + n], + grpo_batch.completion_mask, + device, + response_token_ids=all_rti[offset:offset + n]) + grpo_batch.teacher_per_token_logps = teacher_out.topk_logprobs[..., 0] + offset += n + + def _log_teacher_kl_metric(self, mini_batch_data: List[Dict[str, Any]]) -> None: + """OPD-RL: log the per-token teacher KL (k3) averaged over response tokens (monitoring only). + + Note the deliberate asymmetry: the advantage uses the *signed* k1 log-ratio + (``teacher_logp - student_logp``, see ``expand_advantage_to_per_token``), but monitoring uses + the non-negative k3 estimator because it is the better "distance from the teacher" gauge -- + it should decrease over training. They measure different things on purpose.""" + mode = 'train' if self.unwrapped_models[0].training else 'eval' + kl_sum, tok_sum = 0.0, 0.0 + for data in mini_batch_data: + grpo_batch: GRPOBatch = data['grpo_batch'] + k3 = compute_teacher_kl_per_token(grpo_batch.teacher_per_token_logps, grpo_batch.old_per_token_logps, + grpo_batch.completion_mask) + kl_sum += k3.sum().item() + tok_sum += grpo_batch.completion_mask.sum().item() + self._metrics[mode]['teacher_kl'].append(kl_sum / max(tok_sum, 1.0)) + + def _compute_kl_from_batches(self, mini_batch_data: List[Dict[str, Any]]) -> torch.Tensor: + """ + Compute per-sample KL divergence from encoded batches for kl_in_reward. + + The KL is computed as: sum over tokens of (old_logp - ref_logp) for each sample. + + Args: + mini_batch_data: List of encoded batch dictionaries containing: + - old_per_token_logps: [batch_size, max_seq_len] in batch format + - ref_per_token_logps: [batch_size, max_seq_len] in batch format + - completion_mask: [batch_size, max_seq_len] mask for completion tokens + + Returns: + kl_values: Per-sample KL values, shape [total_samples] + """ + kl_list = [] + assert self.beta != 0.0 + + for batch in mini_batch_data: + grpo_batch = batch['grpo_batch'] + # Compute per-token KL: old_logp - ref_logp + per_token_kl = grpo_batch.old_per_token_logps - grpo_batch.ref_per_token_logps # [batch_size, max_seq_len] + + # Compute per-sample KL by summing over completion tokens + sample_kl = (per_token_kl * grpo_batch.completion_mask).sum(-1) # [batch_size] + kl_list.append(sample_kl) + + # Concatenate all KL values and gather across ranks + kl_values = torch.cat(kl_list, dim=0) + kl_values = gather(kl_values) + + return kl_values + + @contextmanager + def _disable_maxlength_template_context(self, template: Template): + # The max_length for prompt and completion has already been restricted, so there is no need for max_length here. + max_length = template.max_length + template.max_length = None + try: + yield + finally: + template.max_length = max_length + + @property + def on_policy(self): + return self.steps_per_generation == 1 + + @profiling_decorator + def forward_step(self, data_iterator, model): + args = self.args + data = next(data_iterator) + grpo_batch: GRPOBatch = data.pop('grpo_batch') + data = self._prepare_batch(data) + data.pop('loss_scale', None) + + if self.enable_routing_replay and RouterReplayHelper.is_replay_backward_action(model.config): + router_instance_list = RouterReplayHelper.get_micro_batch_router_list(model.config) + for router in router_instance_list: + router.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD) + if self.enable_routing_replay and RouterReplayHelper.is_replay_forward_action(model.config): + layers_topk_idx = data.pop('routed_experts', None) + set_router_replay_data(layers_topk_idx, model.config) + + labels = data.pop('labels', None) + packed_seq_params = data.get('packed_seq_params') + max_seq_len = grpo_batch.completion_mask.shape[1] + micro_batch_size = self.micro_batch_size + + # data is now clean model forward kwargs (template.encode guarantees this; + # grpo_batch / loss_scale / routed_experts / labels all popped above) + is_pp_last_stage = mpu.is_pipeline_last_stage() + output_tensor = model(**data) + if is_pp_last_stage and output_tensor is not None: + logits_packed = output_tensor + if self.temperature != 1.0: + logits_packed.div_(self.temperature) + per_token_logps_packed, per_token_entropy_packed = compute_logps_and_entropy_from_logits( + logits_packed, labels, compute_entropy=self.compute_entropy) + + if args.context_parallel_size > 1: + num_samples = packed_seq_params.seq_lens.shape[0] if args.padding_free else micro_batch_size + cp_size = args.context_parallel_size + per_token_logps_packed = reconstruct_tensor_cp(cp_size, per_token_logps_packed, packed_seq_params, + num_samples) + if per_token_entropy_packed is not None: + per_token_entropy_packed = reconstruct_tensor_cp(cp_size, per_token_entropy_packed, + packed_seq_params, num_samples) + + if args.padding_free: + # Pad from rmpad [1, total_tokens] to batch format [batch_size, max_seq_len] + per_token_logps, _ = pad_logps_back_to_batch( + logps_rmpad=per_token_logps_packed, + logits_to_keep=max_seq_len, + batch_size=micro_batch_size, + seq_lengths=grpo_batch.seq_lengths) + if per_token_entropy_packed is not None: + per_token_entropy, _ = pad_logps_back_to_batch( + logps_rmpad=per_token_entropy_packed, + logits_to_keep=max_seq_len, + batch_size=micro_batch_size, + seq_lengths=grpo_batch.seq_lengths, + pad_value=float('nan')) + else: + per_token_entropy = None + else: + per_token_logps = per_token_logps_packed + per_token_entropy = per_token_entropy_packed + + output_tensor = per_token_logps + data['per_token_entropy'] = per_token_entropy + + if self.enable_routing_replay and RouterReplayHelper.is_replay_forward_action(model.config): + router_instance_list = RouterReplayHelper.get_micro_batch_router_list(model.config) + for router in router_instance_list: + router.set_router_replay_action(RouterReplayAction.REPLAY_BACKWARD) + + data['grpo_batch'] = grpo_batch + data['labels'] = labels + return output_tensor, partial(self.loss_func, data=data) + + @profiling_decorator + def loss_func(self, output_tensor: torch.Tensor, data: Dict[str, Any]): + grpo_batch: GRPOBatch = data['grpo_batch'] + # Get pre-padded data in batch format [batch_size, max_seq_len] + advantages = grpo_batch.advantages # [batch_size, max_seq_len] (per-token, expanded at batch construction) + completion_mask = grpo_batch.completion_mask # [batch_size, max_seq_len] + truncated_mask = grpo_batch.truncated_mask # [batch_size] + micro_batch_size = self.micro_batch_size + + # Get pre-computed per_token_logps and per_token_entropy from forward_step + # These are already in batch format [batch_size, max_seq_len] + per_token_logps = output_tensor + per_token_entropy = data.get('per_token_entropy') + + # Get pre-padded ref/old/rollout logps from grpo_batch + ref_per_token_logps = grpo_batch.ref_per_token_logps # [batch_size, max_seq_len] or None + old_per_token_logps = grpo_batch.old_per_token_logps # [batch_size, max_seq_len] + rollout_per_token_logps = grpo_batch.rollout_per_token_logps # [batch_size, max_seq_len] or None + + # Rollout importance sampling correction + rollout_correction_metrics = {} + rollout_is_weights = None + should_compute_rollout_metrics = (not self.disable_rollout_importance_sampling + and (self.rollout_importance_sampling_mode is not None + or self.log_rollout_offpolicy_metrics)) + if should_compute_rollout_metrics: + dp_group = mpu.get_data_parallel_group(with_context_parallel=True) + has_flag = torch.tensor([1 if rollout_per_token_logps is not None else 0], + dtype=torch.int32, + device=per_token_logps.device) + torch.distributed.all_reduce(has_flag, op=torch.distributed.ReduceOp.MIN, group=dp_group) + should_compute_rollout_metrics = has_flag.item() > 0 + if should_compute_rollout_metrics: + # Compute off-policy diagnostic metrics + rollout_correction_metrics = self._compute_rollout_offpolicy_metrics(old_per_token_logps, + rollout_per_token_logps, + completion_mask) + + # Apply importance sampling correction if mode is enabled + if self.rollout_importance_sampling_mode is not None: + rollout_log_ratio = old_per_token_logps - rollout_per_token_logps + rollout_is_weights = self._apply_rollout_importance_sampling(rollout_log_ratio, completion_mask) + + # Compute IS-specific metrics + is_metrics = self._compute_is_correction_metrics(rollout_log_ratio, rollout_is_weights, completion_mask) + rollout_correction_metrics.update(is_metrics) + + # Compute completion lengths before any mask modifications for accurate logging + completion_lengths = completion_mask.sum(dim=-1) + + # Apply truncated_mask filter (now in batch format) + if self.args.overlong_filter and truncated_mask.any(): + if truncated_mask.all(): + logger.warning('All completions are truncated in this batch. Loss and grad_norm will be 0. ' + 'Consider increasing max_completion_length') + # Expand truncated_mask from [batch_size] to [batch_size, max_seq_len] + truncated_mask_expanded = truncated_mask.unsqueeze(-1).expand_as(completion_mask) + completion_mask = completion_mask & (~truncated_mask_expanded) + + # Compute KL divergence if needed + # Only compute KL for loss if kl_in_reward=False (GRPO style) + # When kl_in_reward=True, KL penalty is already applied to rewards in _compute_advantages + if self.beta != 0.0 and ref_per_token_logps is not None and not self.kl_in_reward: + safe_ratio = torch.clamp(ref_per_token_logps - per_token_logps, min=-20, max=20) + per_token_kl = torch.clamp(torch.exp(safe_ratio) - safe_ratio - 1, min=-10, max=10) + else: + per_token_kl = None + + # Compute log ratio for importance sampling + log_ratio = per_token_logps - old_per_token_logps + + # Compute importance weights based on level + if self.importance_sampling_level == 'token': + log_importance_weights = log_ratio + elif self.importance_sampling_level in ['sequence', 'sequence_token']: + # Sequence-level: compute mean log ratio per sequence + seq_level_log_weights = ((log_ratio * completion_mask).sum(-1) + / completion_mask.sum(-1).clamp(min=1.0)).unsqueeze(-1) + if self.importance_sampling_level == 'sequence': + log_importance_weights = seq_level_log_weights + else: + seq_level_log_weight = seq_level_log_weights.detach() + log_importance_weights = per_token_logps - per_token_logps.detach() + seq_level_log_weight + else: + raise ValueError( + f"Unknown importance sampling level: {self.importance_sampling_level}. Possible values are 'token' " + ",'sequence' and 'sequence_token'.") + + coef_1 = torch.exp(log_importance_weights) + + # advantages is per-token [B, T] (expanded at batch construction so the OPD-RL signed teacher + # log-ratio is added per token). Edge loss types that need a per-sequence advantage (real / fipo / + # off_policy_sequence_mask) are not supported with a teacher. + if self._has_teacher and (self.loss_type in ['real', 'fipo'] + or self.off_policy_sequence_mask_delta is not None): + raise ValueError(f'OPD-RL (teacher) does not support loss_type={self.loss_type!r} / ' + 'off_policy_sequence_mask.') + + fipo_metrics = None + if self.loss_type == 'cispo': + clamped_ratios = torch.clamp(coef_1, max=self.epsilon_high).detach() + per_token_loss = -clamped_ratios * advantages * per_token_logps + elif self.loss_type == 'sapo': + gate_pos = torch.sigmoid(self.tau_pos * (coef_1 - 1)) * (4.0 / self.tau_pos) + gate_neg = torch.sigmoid(self.tau_neg * (coef_1 - 1)) * (4.0 / self.tau_neg) + is_positive = advantages > 0 + soft_gate = torch.where(is_positive, gate_pos, gate_neg) + per_token_loss = -soft_gate * advantages + elif self.loss_type in ['grpo', 'bnpo', 'dr_grpo', 'dapo', 'fipo']: + if self.loss_type == 'fipo': + fipo_weight, fipo_metrics = self._compute_fipo_influence(log_ratio, coef_1, advantages, completion_mask) + + coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high) + if self.args.delta is not None: + coef_1 = torch.clamp(coef_1, max=self.args.delta) + + per_token_loss1 = coef_1 * advantages + per_token_loss2 = coef_2 * advantages + per_token_loss = -torch.min(per_token_loss1, per_token_loss2) + if self.loss_type == 'fipo': + per_token_loss = per_token_loss * fipo_weight + elif self.loss_type == 'real': + per_token_loss = torch.zeros_like(per_token_logps) + else: + raise ValueError(f'Unknown loss type: {self.loss_type}') + + # Compute and apply entropy mask if enabled + # entropy_mask zeros out gradients for low-entropy (confident) tokens + entropy_mask = None + entropy_metrics = {} + if self.compute_entropy and per_token_entropy is not None: + # Fill padded tokens with NaN for correct quantile computation + entropies = per_token_entropy.masked_fill(completion_mask == 0, float('nan')) + + if self.log_entropy: + # Log entropy statistics + per_completion_entropies_mean = torch.nanmean(entropies, dim=1) + global_entropies_mean = gather( + per_completion_entropies_mean, group=mpu.get_data_parallel_group(with_context_parallel=True)) + valid_mask = ~torch.isnan(global_entropies_mean) + entropy_metrics = { + 'entropy_mean': + global_entropies_mean.nanmean(), + 'entropy_max': + global_entropies_mean[valid_mask].max() if valid_mask.any() else torch.tensor( + float('nan'), device=global_entropies_mean.device), + 'entropy_min': + global_entropies_mean[valid_mask].min() if valid_mask.any() else torch.tensor( + float('nan'), device=global_entropies_mean.device), + } + + # Compute entropy threshold and mask for top_entropy_quantile + if self.top_entropy_quantile < 1.0: + # Compute threshold across all completion tokens in the batch + entropy_threshold = torch.nanquantile(entropies.flatten().float(), 1 - self.top_entropy_quantile) + entropy_metrics['entropy_threshold'] = entropy_threshold + # Only keep tokens with entropy >= threshold (high uncertainty tokens) + entropy_mask = entropies >= entropy_threshold + # Apply entropy mask to per_token_loss + per_token_loss = per_token_loss * entropy_mask + + # Add KL penalty if needed + if self.beta != 0.0 and per_token_kl is not None: + per_token_loss = per_token_loss + self.beta * per_token_kl + + # Apply rollout importance sampling weights if available + if rollout_is_weights is not None and self.rollout_importance_sampling_mode is not None: + per_token_loss = per_token_loss * rollout_is_weights + + # Apply off-policy sequence masking if enabled + # Mask out sequences where delta > threshold AND advantage < 0 + if self.off_policy_sequence_mask_delta is not None: + old_policy_per_token_logps = rollout_per_token_logps if rollout_per_token_logps is not None \ + else old_per_token_logps + # advantages is per-token [B, T]; the mask needs a per-sequence scalar. + seq_advantages = (advantages * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0) + off_policy_seq_mask = self._compute_off_policy_sequence_mask(per_token_logps, old_policy_per_token_logps, + completion_mask, seq_advantages) + # Expand sequence mask to token level and apply to completion_mask + off_policy_seq_mask_expanded = off_policy_seq_mask.unsqueeze(-1).expand_as(completion_mask) + completion_mask = completion_mask & off_policy_seq_mask_expanded + + if self.loss_type in ['grpo', 'sapo']: + # Per-sample mean, then batch mean + loss = ((per_token_loss * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)).mean() + elif self.loss_type == 'bnpo': + loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + elif self.loss_type == 'dr_grpo': + loss = (per_token_loss * completion_mask).sum() / (micro_batch_size * self.max_completion_length) + elif self.loss_type in ['cispo', 'dapo', 'fipo']: + # CISPO, DAPO, and FIPO: Normalize by total completion tokens across all processes + num_items_in_batch = grpo_batch.num_items_in_batch + dp_size = mpu.get_data_parallel_world_size() + normalizer = num_items_in_batch / dp_size + loss = (per_token_loss * completion_mask).sum() / normalizer.clamp(min=1.0) + elif self.loss_type == 'real': + global_scores = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0) + + group_scores = global_scores.view(-1, self.num_generations) + # advantages is per-token [B, T] (constant across tokens without a teacher); reduce + # to a per-sequence scalar for the group pos/neg split. + seq_advantages = (advantages * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0) + group_rewards = seq_advantages.view(-1, self.num_generations) + + pos_mask = (group_rewards > 0) + neg_mask = (group_rewards <= 0) + valid_mask = (pos_mask.sum(dim=1) != 0) & (neg_mask.sum(dim=1) != 0) + + if not valid_mask.any(): + loss = torch.tensor(0., device=global_scores.device) * global_scores.mean() + else: + batch_scores = group_scores[valid_mask] + batch_pos_mask = pos_mask[valid_mask] + batch_neg_mask = neg_mask[valid_mask] + + scaled_scores = batch_scores / self.real_tau + zeros = torch.zeros(batch_scores.size(0), 1, device=batch_scores.device, dtype=batch_scores.dtype) + + # Negative Loss: log(1 + sum(e^{S_neg})) + neg_input = scaled_scores.masked_fill(~batch_neg_mask, float('-inf')) + neg_loss = torch.logsumexp(torch.cat([neg_input, zeros], dim=1), dim=1) + + # Positive Loss: log(1 + sum(e^{-S_pos})) + pos_input = (-scaled_scores).masked_fill(~batch_pos_mask, float('-inf')) + pos_loss = torch.logsumexp(torch.cat([pos_input, zeros], dim=1), dim=1) + + loss = (neg_loss + pos_loss).sum() / group_rewards.size(0) + + if self.beta != 0.0 and per_token_kl is not None: + kl_loss = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + loss = loss + kl_loss * self.beta + else: + raise ValueError(f'Unknown loss type: {self.loss_type}') + + avg_metric = { + 'loss': loss.clone().detach(), + } + custom_metrics = {} + total_completion_lengths = gather( + completion_lengths, group=mpu.get_data_parallel_group(with_context_parallel=True)) + custom_metrics = { + 'completions/mean_length': total_completion_lengths.float().mean(), + 'completions/max_length': total_completion_lengths.float().max(), + 'completions/min_length': total_completion_lengths.float().min(), + } + + # Add entropy metrics to custom_metrics + if entropy_metrics: + for key, value in entropy_metrics.items(): + if isinstance(value, torch.Tensor): + custom_metrics[f'entropy/{key}'] = value.clone().detach() + else: + custom_metrics[f'entropy/{key}'] = torch.tensor(value, device=loss.device) + + if self.beta != 0.0 and per_token_kl is not None: + kl_value = (per_token_kl * completion_mask).sum() / completion_mask.sum().clamp(min=1.0) + avg_metric['kl'] = kl_value.clone().detach() + + mode = 'train' if self.unwrapped_models[0].training else 'eval' + + # Compute clipping metrics + completion_token_count = completion_mask.sum().clamp(min=1.0) + if fipo_metrics is not None: + avg_metric['fipo/future_kl_mean'] = ((fipo_metrics['future_kl'] * completion_mask).sum() + / completion_token_count).clone().detach() + avg_metric['fipo/influence_weight_mean'] = ((fipo_metrics['influence_weight'] * completion_mask).sum() + / completion_token_count).clone().detach() + avg_metric['fipo/safety_keep_ratio'] = ((fipo_metrics['safety_mask'].float() * completion_mask).sum() + / completion_token_count).clone().detach() + + if self.loss_type == 'cispo': + # CISPO: Only track upper bound clipping + # coef_1 is [batch_size, max_seq_len] or [batch_size, 1] depending on importance_sampling_level + is_cispo_clipped = (coef_1 > self.epsilon_high) & (advantages > 0) + cispo_clip_ratio = (is_cispo_clipped.float() * completion_mask).sum() / completion_token_count + # Store local clip ratio, _all_reduce_metric will handle averaging across ranks + self._metrics[mode]['cispo_clip_ratio'].append(cispo_clip_ratio) + elif self.loss_type in ['sapo', 'real']: + # SAPO / REAL: No hard clipping, skip clipping metrics + pass + elif self.loss_type in ['grpo', 'bnpo', 'dr_grpo', 'dapo', 'fipo']: + # coef_1 is [batch_size, max_seq_len] or [batch_size, 1] depending on importance_sampling_level + # Use exp(log_importance_weights) to get the original ratios before clamping + coef_1_for_metrics = torch.exp(log_importance_weights) + is_low_clipped = (coef_1_for_metrics < 1 - self.epsilon_low) & (advantages < 0) + is_high_clipped = (coef_1_for_metrics > 1 + self.epsilon_high) & (advantages > 0) + low_clip = (is_low_clipped.float() * completion_mask).sum() / completion_token_count + high_clip = (is_high_clipped.float() * completion_mask).sum() / completion_token_count + is_region_clipped = is_low_clipped | is_high_clipped + clip_ratio = (is_region_clipped.float() * completion_mask).sum() / completion_token_count + + # For min/max, we need to gather values from all ranks to compute global min/max + # For mean, let _all_reduce_metric handle averaging + gathered_low_clip = gather( + low_clip.unsqueeze(0), group=mpu.get_data_parallel_group(with_context_parallel=True)) + gathered_high_clip = gather( + high_clip.unsqueeze(0), group=mpu.get_data_parallel_group(with_context_parallel=True)) + + # Store local values for mean (will be averaged by _all_reduce_metric) + self._metrics[mode]['clip_ratio/low_mean'].append(low_clip.item()) + self._metrics[mode]['clip_ratio/high_mean'].append(high_clip.item()) + self._metrics[mode]['clip_ratio/region_mean'].append(clip_ratio.item()) + # Store global min/max in custom_metrics (not through _all_reduce_metric to avoid incorrect averaging) + custom_metrics['clip_ratio/low_min'] = gathered_low_clip.min() + custom_metrics['clip_ratio/high_max'] = gathered_high_clip.max() + + # Add rollout correction metrics + if rollout_correction_metrics: + for key, value in rollout_correction_metrics.items(): + if isinstance(value, torch.Tensor): + custom_metrics[f'rollout_correction/{key}'] = value.clone().detach() + else: + custom_metrics[f'rollout_correction/{key}'] = torch.tensor(value, device=loss.device) + + if self._metrics[mode]: + addition_metrics = { + key: torch.tensor(sum(val) / len(val), device=loss.device) + for key, val in self._metrics[mode].items() + } + avg_metric.update(addition_metrics) + # Reward metrics are produced once per rollout, while the same rollout can be + # consumed by multiple optimization steps. Once handed to Megatron's logging + # aggregator, clear them so they are neither reported repeatedly nor averaged + # together with every rollout seen since the start of training. + self._metrics[mode].clear() + + avg_metric = self._all_reduce_metric(avg_metric) + reporting_metric = {**avg_metric, **custom_metrics} + + if (self._step - 1) % self.steps_per_generation == 0: + self._flush_log_completions() + + return loss, reporting_metric + + @profiling_decorator + def resample_encode_failed_inputs(self, inputs: DataType, max_resample_rounds: int = 10) -> DataType: + """ + Attempt to encode each input using the template. If encoding fails, + resample from a backup iterator until we have enough valid samples. + + This method handles two cases: + 1. Prompt length exceeds max_length + 2. Encoding failures (e.g., multimodal data processing errors) + + Args: + inputs (DataType): A list of input data samples, each containing a `messages` field. + max_resample_rounds (int, optional): Maximum number of resample rounds. + Each round processes samples from pending_samples buffer. Defaults to 10. + + Returns: + DataType: A list of successfully encoded input samples with the same length as inputs. + + Raises: + RuntimeError: If we cannot collect enough valid samples after max_resample_rounds. + """ + return resample_encode_failed_inputs( + self.template, + self.resample_data_iterator, + inputs, + max_resample_rounds=max_resample_rounds, + strip_response=True, + ) + + def get_num_iters_per_step(self): + mode = 'train' if self.unwrapped_models[0].training else 'eval' + cache_key = f'_num_iters_per_step_{mode}' + if hasattr(self, cache_key): + return getattr(self, cache_key) + dp_size = mpu.get_data_parallel_world_size() + completions_batch_size = self.generation_batch_size if mode == 'train' else self.global_batch_size + completions_to_rollout = completions_batch_size // dp_size + # completions will be repeated num_generations times after + # so we need to divide num_iters_per_step by num_generations to get prompt batch size + num_generations = self.num_generations if mode == 'train' else self.num_generations_eval + prompts_to_rollout = completions_to_rollout // num_generations + # every iter will generate micro_batch_size prompts + num_iters_per_step = prompts_to_rollout // self.micro_batch_size + assert num_iters_per_step > 0, ( + f'num_iters_per_step={num_iters_per_step} <= 0. ' + f'This means no prompts will be generated' + f'generation_batch_size={self.generation_batch_size}, ' + f'data_parallel_world_size={mpu.get_data_parallel_world_size()}, ' + f'num_generations={num_generations}, ' + f'micro_batch_size={self.micro_batch_size}. ' + 'Please adjust these parameters so that ' + 'generation_batch_size // data_parallel_world_size // num_generations // micro_batch_size >= 1.') + setattr(self, cache_key, num_iters_per_step) + return num_iters_per_step + + def get_local_rollout_batch(self, batch: List[GRPOSample]) -> List[GRPOSample]: + mode = 'train' if self.unwrapped_models[0].training else 'eval' + num_generations = self.num_generations if mode == 'train' else self.num_generations_eval + # repeat num_generations times + rollout_group = self._get_rollout_group() + global_rollout_batch = [deepcopy(item) for item in batch for _ in range(num_generations)] + # get local rollout data + rollout_rank = torch.distributed.get_rank(group=rollout_group) + rollout_group_size = torch.distributed.get_world_size(group=rollout_group) + + per_device_batch_size = self.per_device_generation_batch_size + if mode != 'train': + per_device_batch_size //= self.steps_per_generation + assert rollout_group_size * per_device_batch_size == len(global_rollout_batch) + data_slice = slice(rollout_rank * per_device_batch_size, (rollout_rank + 1) * per_device_batch_size) + rollout_batch = global_rollout_batch[data_slice] + return rollout_batch + + @contextmanager + def _template_context(self, template: Template): + # The max_length for prompt and completion has already been restricted, so there is no need for max_length here. + max_length = template.max_length + template.max_length = None + try: + yield + finally: + template.max_length = max_length + + def _prepare_metrics(self): + # Shared logging infrastructure (prompt/completion/jsonl/wandb/swanlab) + self._prepare_logging() + # GRPO-specific: add rewards and advantages columns + self._logs['rewards'] = defaultdict(lambda: deque(maxlen=self.args.generation_batch_size)) + self._logs['advantages'] = deque(maxlen=self.args.generation_batch_size) + + self.init_custom_metric = False + self._metrics = {'train': defaultdict(list), 'eval': defaultdict(list)} + + def _compute_sequence_level_ratios(self, is_ratio: torch.Tensor, completion_mask: torch.Tensor) -> torch.Tensor: + """ + Helper function to compute sequence-level importance sampling ratios. + + Args: + is_ratio: Token-level IS ratios, shape [batch_size, seq_len] + completion_mask: Boolean mask for completion tokens, shape [batch_size, seq_len] + + Returns: + Sequence-level ratios as geometric mean of token-level ratios, shape [batch_size] + """ + log_ratio = torch.log(is_ratio.clamp(min=1e-10)) + # Compute per-sequence mean of log ratios + seq_log_ratios = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0) + seq_ratios = torch.exp(seq_log_ratios) + return seq_ratios + + def _apply_rollout_importance_sampling(self, rollout_log_ratio: torch.Tensor, + completion_mask: torch.Tensor) -> torch.Tensor: + """ + Apply rollout importance sampling correction using one of four modes. + + Args: + rollout_log_ratio: log(π_θ / π_rollout) per token, shape [batch_size, seq_len] + completion_mask: Boolean mask for completion tokens, shape [batch_size, seq_len] + + Returns: + IS weights to multiply with loss, same shape as rollout_log_ratio + """ + mode = self.rollout_importance_sampling_mode + threshold = self.rollout_importance_sampling_threshold + + # Clamp log_ratio to prevent numerical overflow from padding values + SAFETY_BOUND = 20.0 + rollout_log_ratio_safe = torch.clamp(rollout_log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND) + + # Compute importance sampling ratios: exp(log_ratio) + is_ratio = torch.exp(rollout_log_ratio_safe) + + if mode == 'token_truncate': + # Token-level truncated IS: clip ratios from above at threshold + is_weights = torch.clamp(is_ratio, max=threshold) + + elif mode == 'token_mask': + # Token-level masked IS: mask out tokens with ratio > threshold + is_weights = torch.where(is_ratio <= threshold, is_ratio, torch.zeros_like(is_ratio)) + + elif mode == 'sequence_truncate': + # Sequence-level truncated IS: compute sequence-level ratio and clip + seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask) + clipped_seq_ratios = torch.clamp(seq_ratios, max=threshold) + + # Expand back to token-level [batch_size] -> [batch_size, seq_len] + is_weights = clipped_seq_ratios.unsqueeze(-1).expand_as(is_ratio) + + elif mode == 'sequence_mask': + # Sequence-level masked IS: mask entire sequences with ratio > threshold + seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask) + seq_mask = (seq_ratios <= threshold).float() + + # Expand mask to token-level and apply to original token-level ratios + seq_mask_expanded = seq_mask.unsqueeze(-1).expand_as(is_ratio) + is_weights = is_ratio * seq_mask_expanded + else: + return is_ratio + + return is_weights + + def _compute_off_policy_sequence_mask( + self, + per_token_logps: torch.Tensor, + old_policy_per_token_logps: torch.Tensor, + completion_mask: torch.Tensor, + advantages: torch.Tensor, + ) -> torch.Tensor: + """ + Compute off-policy sequence mask to filter out sequences that deviate too much + from the old/rollout policy AND have negative advantage. + + This implements the Off-Policy Sequence Masking technique from DeepSeek-V3.2 + (https://arxiv.org/abs/2512.02556). The mask filters sequences where: + 1. mean(old_policy_logps - policy_logps) > off_policy_sequence_mask_delta + 2. AND advantage < 0 + + Args: + per_token_logps: Log probs from current policy, shape [batch_size, seq_len] + old_policy_per_token_logps: Log probs from old/rollout policy, shape [batch_size, seq_len]. + Uses rollout_per_token_logps if available, otherwise old_per_token_logps. + completion_mask: Boolean mask for completion tokens, shape [batch_size, seq_len] + advantages: Advantage values per sample, shape [batch_size] + + Returns: + Sequence mask, shape [batch_size], True = keep sequence, False = mask out + """ + # Compute per-token log ratio: log(π_old / π_current) + # Following DeepSeek-V3.2: positive delta means old policy assigns higher prob + log_ratio = old_policy_per_token_logps - per_token_logps + + # Compute sequence-level mean of log ratio + seq_mean_log_ratio = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0) + + # Mask condition: delta > threshold AND advantage < 0 + # Keep sequences that do NOT meet this condition + exceeds_threshold = seq_mean_log_ratio > self.off_policy_sequence_mask_delta + negative_advantage = advantages < 0 + should_mask = exceeds_threshold & negative_advantage + + # Return mask: True = keep, False = mask out + return ~should_mask + + def _compute_rollout_offpolicy_metrics( + self, + per_token_logps: torch.Tensor, + rollout_per_token_logps: torch.Tensor, + completion_mask: torch.Tensor, + ) -> Dict[str, float]: + """ + Compute off-policy diagnostic metrics (always computed for monitoring). + + These metrics help diagnose the off-policy gap between rollout and training policies. + + Args: + per_token_logps: Log probs from training policy model, shape [batch_size, seq_len] + rollout_per_token_logps: Log probs from rollout policy, shape [batch_size, seq_len] + completion_mask: Boolean mask for completion tokens, shape [batch_size, seq_len] + + Returns: + Dictionary with off-policy diagnostic metrics + """ + SAFETY_BOUND = 20.0 + metrics = {} + dp_group = mpu.get_data_parallel_group(with_context_parallel=True) + + # Helper function for masked mean + def masked_mean(x, mask, axis=None): + if axis is None: + return (x * mask).sum() / mask.sum().clamp(min=1.0) + else: + return (x * mask).sum(axis) / mask.sum(axis).clamp(min=1.0) + + # 1. Training policy perplexity + # Compute per-sequence mean log prob + mean_log_prob_training = masked_mean(per_token_logps, completion_mask, axis=-1) # [batch_size] + training_ppl = torch.exp(-mean_log_prob_training).mean() + gathered_training_ppl = gather(training_ppl.unsqueeze(0), group=dp_group) + metrics['training_ppl'] = gathered_training_ppl.nanmean() + metrics['training_log_ppl'] = gather((-mean_log_prob_training).mean().unsqueeze(0), group=dp_group).nanmean() + + # 2. Compute rollout off-policy metrics + log_ratio = per_token_logps - rollout_per_token_logps + log_ratio = log_ratio * completion_mask + + # 2a. kl: Direct estimator for KL(π_rollout || π_training) + kl = masked_mean(-log_ratio, completion_mask) + metrics['kl'] = gather(kl.unsqueeze(0), group=dp_group).nanmean() + + # 2b. k3_kl: K3 estimator for KL + log_ratio_safe = torch.clamp(log_ratio, min=-20, max=20) + k3_kl_matrix = torch.clamp(torch.exp(log_ratio_safe) - log_ratio_safe - 1, min=-10, max=10) + k3_kl = masked_mean(k3_kl_matrix, completion_mask) + metrics['k3_kl'] = gather(k3_kl.unsqueeze(0), group=dp_group).nanmean() + + # 2c. Rollout policy perplexity + mean_log_prob_rollout = masked_mean(rollout_per_token_logps, completion_mask, axis=-1) # [batch_size] + rollout_ppl = torch.exp(-mean_log_prob_rollout).mean() + metrics['rollout_ppl'] = gather(rollout_ppl.unsqueeze(0), group=dp_group).nanmean() + metrics['rollout_log_ppl'] = gather((-mean_log_prob_rollout).mean().unsqueeze(0), group=dp_group).nanmean() + + # 2d. Log PPL difference + log_ppl_diff = mean_log_prob_rollout - mean_log_prob_training + metrics['log_ppl_diff'] = gather(log_ppl_diff.mean().unsqueeze(0), group=dp_group).nanmean() + metrics['log_ppl_abs_diff'] = gather(log_ppl_diff.abs().mean().unsqueeze(0), group=dp_group).nanmean() + metrics['log_ppl_diff_max'] = gather(log_ppl_diff.max().unsqueeze(0), group=dp_group).max() + metrics['log_ppl_diff_min'] = gather(log_ppl_diff.min().unsqueeze(0), group=dp_group).min() + + # 2e. PPL ratio + ppl_ratio = torch.exp(log_ppl_diff).mean() + metrics['ppl_ratio'] = gather(ppl_ratio.unsqueeze(0), group=dp_group).nanmean() + + # 2f. Chi-squared divergence + log_ratio_safe = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND) + rho_token = torch.exp(log_ratio_safe) + rho_squared_token = rho_token.square() + chi2_token = masked_mean(rho_squared_token, completion_mask) - 1.0 + metrics['chi2_token'] = gather(chi2_token.unsqueeze(0), group=dp_group).nanmean() + + # Sequence-level chi2 + log_ratio_mean = masked_mean(log_ratio, completion_mask, axis=-1) # [batch_size] + log_ratio_mean_safe = torch.clamp(log_ratio_mean, min=-SAFETY_BOUND, max=SAFETY_BOUND) + rho_geo = torch.exp(log_ratio_mean_safe) + chi2_seq = (rho_geo.square().mean() - 1.0) + metrics['chi2_seq'] = gather(chi2_seq.unsqueeze(0), group=dp_group).nanmean() + + return metrics + + def _compute_is_correction_metrics( + self, + rollout_log_ratio: torch.Tensor, + is_weights: torch.Tensor, + completion_mask: torch.Tensor, + ) -> Dict[str, float]: + """ + Compute importance sampling correction metrics (ess, clipped_frac, is_weight_mean). + Only called when rollout_importance_sampling_mode is enabled. + + Args: + rollout_log_ratio: Log ratio log(π_policy / π_rollout), shape [batch_size, seq_len] + is_weights: Importance sampling weights after correction, shape [batch_size, seq_len] + completion_mask: Boolean mask for completion tokens, shape [batch_size, seq_len] + + Returns: + Dictionary with IS-specific metrics + """ + metrics = {} + SAFETY_BOUND = 20.0 + threshold = self.rollout_importance_sampling_threshold + threshold_lower = 1.0 / threshold + dp_group = mpu.get_data_parallel_group(with_context_parallel=True) + + # Helper function for masked mean + def masked_mean(x, mask): + return (x * mask).sum() / mask.sum().clamp(min=1.0) + + # Compute IS ratio with safety bounds + log_ratio_safe = torch.clamp(rollout_log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND) + is_ratio = torch.exp(log_ratio_safe) + + # 1. IS weight statistics + mean_is_weight = masked_mean(is_weights, completion_mask) + metrics['is_weight_mean'] = gather(mean_is_weight.unsqueeze(0), group=dp_group).nanmean() + + # 2. Compute Effective Sample Size (ESS) for IS weights + weights_for_ess = is_weights.clamp(min=threshold_lower, max=threshold) + mean_for_ess = masked_mean(weights_for_ess, completion_mask) + is_weights_normalized = weights_for_ess / (mean_for_ess + 1e-8) + ess = 1.0 / masked_mean(is_weights_normalized.square(), completion_mask).clamp(min=1e-10) + metrics['ess'] = gather(ess.unsqueeze(0), group=dp_group).nanmean() + + # 3. Fraction of clipped/masked samples + if self.rollout_importance_sampling_mode in ['token_truncate', 'token_mask']: + # Token-level + if self.rollout_importance_sampling_mode == 'token_truncate': + clipped_frac = masked_mean((is_ratio > threshold).float(), completion_mask) + else: # token_mask + clipped_frac = masked_mean((is_weights == 0).float(), completion_mask) + metrics['clipped_frac'] = gather(clipped_frac.unsqueeze(0), group=dp_group).nanmean() + else: + # Sequence-level (both truncate and mask) + seq_ratios = self._compute_sequence_level_ratios(is_ratio, completion_mask) + clipped_frac = (seq_ratios > threshold).float().mean() + metrics['clipped_frac'] = gather(clipped_frac.unsqueeze(0), group=dp_group).nanmean() + + return metrics + + def _collect_config_info(self) -> Dict[str, str]: + config = { + 'dynamic_sample': str(self.args.dynamic_sample), + 'importance_sampling_level': str(self.args.importance_sampling_level), + 'advantage_estimator': str(self.args.advantage_estimator), + 'offpolicy_sequence_mask': 'enable' if self.args.off_policy_sequence_mask_delta is not None else 'disable', + 'rollout_importance_sampling': + 'enable' if self.args.rollout_importance_sampling_mode is not None else 'disable', + 'loss_type': str(self.args.loss_type) + } + return config diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/kto_trainer.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/kto_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..f7275af6ff970ba614266b42d2c6cb0132f3d53d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/kto_trainer.py @@ -0,0 +1,165 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from collections import namedtuple +from functools import partial +from megatron.core import mpu +from trl import KTOTrainer +from typing import Any, Dict + +from swift.utils import get_current_device, get_logger +from .rlhf_mixin import MegatronRLHFTrainer + +logger = get_logger() + + +class DummyKTOTrainer(KTOTrainer): + # For reusing the kto_loss function in TRL. + + def gather_for_metrics(self, input_data, *args, **kwargs): + output_tensors = torch.empty( + mpu.get_data_parallel_world_size() * input_data.numel(), + dtype=input_data.dtype, + device=input_data.device, + ) + torch.distributed.all_gather_into_tensor(output_tensors, input_data, group=mpu.get_data_parallel_group()) + return output_tensors + + def __init__(self, args): + self.accelerator = namedtuple('Accelerator', ['device', 'gather_for_metrics'])( + device=get_current_device(), gather_for_metrics=self.gather_for_metrics) + self.loss_type = args.loss_type + self.beta = args.beta + self.desirable_weight = args.desirable_weight + self.undesirable_weight = args.undesirable_weight + self.calculate_KL = args.calculate_KL + + +class MegatronKTOTrainer(MegatronRLHFTrainer): + + def __init__(self, args, template): + super().__init__(args, template) + self.dummy_kto_trainer = DummyKTOTrainer(args) + + def _kto_get_logps(self, output_tensor, data, is_KL: bool, is_ref: bool, length: int): + output = self._get_input_tensor(output_tensor, is_KL, is_ref, length, dim=1) + return self.get_logps(output, data['labels'], data.get('packed_seq_params')) + + def _get_kto_length(self, data: Dict[str, Any]) -> int: + if 'packed_seq_params' in data: + return data['packed_seq_params'].cu_seqlens_q[-1] // self.args.context_parallel_size + else: + return data['position_ids'].shape[-1] + + def loss_func(self, output_tensor, *, data, kl_data, label): + length = self._get_kto_length(data) + policy_logps = self._kto_get_logps(output_tensor, data, False, False, length) + ref_logps = self._kto_get_logps(output_tensor, data, False, True, length) + if self.args.calculate_KL: + policy_KL_logps = self._kto_get_logps(output_tensor, kl_data, True, False, length) + ref_KL_logps = self._kto_get_logps(output_tensor, kl_data, True, True, length) + else: + policy_KL_logps, ref_KL_logps = None, None + label = output_tensor.new_tensor(label, dtype=torch.bool) + policy_chosen_logps = policy_logps[label] + policy_rejected_logps = policy_logps[~label] + ref_chosen_logps = ref_logps[label] + ref_rejected_logps = ref_logps[~label] + + loss, chosen_rewards, rejected_rewards, kl = self.dummy_kto_trainer.kto_loss( + policy_chosen_logps, + policy_rejected_logps, + policy_KL_logps, + ref_chosen_logps, + ref_rejected_logps, + ref_KL_logps, + ) + + loss = loss.mean() + mean_metric = { + 'loss': loss.detach().clone(), + 'kl': kl.squeeze().detach(), + } + metric = self._all_reduce_metric(mean_metric) + chosen_count = chosen_rewards.shape[0] + rejected_count = rejected_rewards.shape[0] + sum_metric = { + 'logps/chosen': loss.new_tensor([policy_chosen_logps.detach().nansum(), chosen_count]), + 'logps/rejected': loss.new_tensor([policy_rejected_logps.detach().nansum(), rejected_count]), + 'rewards/chosen': loss.new_tensor([chosen_rewards.nansum(), chosen_count]), + 'rewards/rejected': loss.new_tensor([rejected_rewards.nansum(), rejected_count]), + } + metric.update(self._all_reduce_metric(sum_metric, torch.distributed.ReduceOp.SUM)) + # fix megatron-lm bug + loss = loss / mpu.get_context_parallel_world_size() + return loss, metric + + @staticmethod + def _get_input_tensor(input_tensor, is_KL: bool, is_ref: bool, length: int, dim: int): + # policy, ref, policy_KL, ref_KL + total_length = input_tensor.shape[dim] + KL_length = (total_length - 2 * length) // 2 + slice_list = [0, length, 2 * length, total_length - KL_length, total_length] + idx = is_KL * 2 + is_ref + slice_ = (slice(None), ) * dim + (slice(slice_list[idx], slice_list[idx + 1]), ) + res = input_tensor[slice_] + if is_KL or is_ref: + res = res.detach() + return res + + def forward_step(self, data_iterator, model): + unwrapped_model = model.module.module + input_tensor = unwrapped_model.get_input_tensor() + vp_stage = unwrapped_model.vp_stage + # not support loss_scale + data, kl_data = self.get_batch(data_iterator, vp_stage) + label = data.pop('label') + data.pop('loss_scale', None) + kl_data.pop('loss_scale', None) + + length = self._get_kto_length(data) + if self.args.sequence_parallel: + length //= mpu.get_tensor_model_parallel_world_size() + with torch.no_grad(), self.null_ref_context() as ref_models: + ref_model = ref_models[vp_stage or 0] + if self.args.calculate_KL: + if input_tensor is not None: + ref_model.set_input_tensor(self._get_input_tensor(input_tensor, True, True, length, 0)) + ref_KL_output_tensor = ref_model(**kl_data) + + if input_tensor is not None: + ref_model.set_input_tensor(self._get_input_tensor(input_tensor, False, True, length, 0)) + ref_output_tensor = ref_model(**data) + + if self.args.calculate_KL: + with torch.no_grad(): + if input_tensor is not None: + unwrapped_model.set_input_tensor(self._get_input_tensor(input_tensor, True, False, length, 0)) + KL_output_tensor = model(**kl_data) + + if input_tensor is not None: + unwrapped_model.set_input_tensor(self._get_input_tensor(input_tensor, False, False, length, 0)) + output_tensor = model(**data) + is_pp_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage) + dim = 1 if is_pp_last_stage else 0 + if self.args.calculate_KL: + res = torch.concat([output_tensor, ref_output_tensor, KL_output_tensor, ref_KL_output_tensor], dim=dim) + else: + res = torch.concat([output_tensor, ref_output_tensor], dim=dim) + return res, partial(self.loss_func, data=data, kl_data=kl_data, label=label) + + def _prepare_batch(self, data, vp_stage=None): + res = [] + for key in ['completion_', 'KL_completion_']: + _data = {k[len(key):]: v for k, v in data.items() if k.startswith(key)} + if not self.args.calculate_KL and key == 'KL_completion_': + _data = {} + else: + _data = super()._prepare_batch(_data, vp_stage) + res.append(_data) + res[0]['label'] = data['label'] + return res + + def _log_callback(self, logs, n_steps): + super()._log_callback(logs, n_steps) + if 'rewards/chosen' in logs and 'rewards/rejected' in logs: + logs['rewards/margins'] = logs['rewards/chosen'] - logs['rewards/rejected'] diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/reranker_trainer.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/reranker_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..eda261ea6435a3edc7112bc9eb914e4bb60f597b --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/reranker_trainer.py @@ -0,0 +1,73 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch.nn +from collections import namedtuple +from functools import partial + +from swift.loss import loss_map +from swift.metrics import eval_metrics_map +from swift.utils import get_logger +from .base import BaseMegatronTrainer + +logger = get_logger() +ModelOutputs = namedtuple('ModelOutputs', ['logits']) + + +class MegatronRerankerTrainer(BaseMegatronTrainer): + + def __init__(self, args, template): + super().__init__(args, template) + self._loss_func = loss_map[args.loss_type](args, self) + self.eval_metrics = eval_metrics_map['reranker'](args, self) + + @staticmethod + def _get_listwise_reranker_preds(logits, labels): + positive_indices = torch.nonzero(labels == 1, as_tuple=False).squeeze(-1).tolist() + positive_indices.append(labels.shape[0]) + preds = [] + for i in range(len(positive_indices) - 1): + start, end = positive_indices[i], positive_indices[i + 1] + preds.append(logits[start:end].argmax()) + preds = torch.stack(preds) + labels = torch.tensor([0] * (len(positive_indices) - 1), device=preds.device) + return preds, labels + + def loss_func(self, + output_tensor: torch.Tensor, + *, + labels: torch.Tensor, + packed_seq_params=None, + attention_mask=None): + training = self.unwrapped_models[0].training + logits = self.get_last_tokens(output_tensor, packed_seq_params, attention_mask) + loss = self._loss_func(ModelOutputs(logits=logits), labels) + args = self.args + logits_detach = logits.detach().squeeze(-1) + if not training: + self.eval_metrics.update(logits_detach, labels) + if args.loss_type == 'listwise_reranker': + preds, labels = self._get_listwise_reranker_preds(logits_detach, labels) + else: + preds = (logits_detach.detach() > 0).long() + acc = (preds == labels).float().mean() + metric = {'loss': loss.detach().clone(), 'acc': acc} + metric = self._all_reduce_metric(metric) + return loss, metric + + def prepare_model(self): + super().prepare_model() + for model in self.unwrapped_models: + lm_model = model.language_model if hasattr(model, 'language_model') else model + lm_model.tokenizer = self.template.tokenizer + + def forward_step(self, data_iterator, model): + vp_stage = model.module.module.vp_stage + data = self.get_batch(data_iterator, vp_stage) + labels = data.pop('labels', None) + output_tensor = model(**data) + loss_func = partial( + self.loss_func, + labels=labels, + packed_seq_params=data.get('packed_seq_params'), + attention_mask=data.get('attention_mask') + if data.get('attention_mask') is not None else data.get('attention_mask_2d')) + return output_tensor, loss_func diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/reward_trainer.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/reward_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..1f7d47d438b6c1740ec4ce6e55ed5a380d941392 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/reward_trainer.py @@ -0,0 +1,46 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from functools import partial +from torch import nn + +from swift.utils import get_logger +from .rlhf_mixin import MegatronRLHFTrainer + +logger = get_logger() + + +class MegatronRewardTrainer(MegatronRLHFTrainer): + + def loss_func(self, output_tensor, *, data): + packed_seq_params = data.get('packed_seq_params') + margin = data.pop('margin', None) + num_samples = output_tensor.shape[0] if packed_seq_params is None else packed_seq_params.seq_lens.shape[0] + rewards = self.get_last_tokens(output_tensor, packed_seq_params, data.get('attention_mask')) + rewards_chosen, rewards_rejected = torch.split(rewards, num_samples // 2, dim=0) + if margin is not None: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected - margin).mean() + else: + loss = -nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean() + if self.args.center_rewards_coefficient is not None: + center_rewards_loss = self.args.center_rewards_coefficient * torch.mean( + (rewards_chosen + rewards_rejected)**2) + loss += center_rewards_loss + rewards_chosen, rewards_rejected = rewards_chosen.detach(), rewards_rejected.detach() + metric = { + 'loss': loss.detach().clone(), + 'rewards/chosen': rewards_chosen.mean(), + 'rewards/rejected': rewards_rejected.mean(), + 'rewards/accuracies': (rewards_chosen > rewards_rejected).float().mean(), + 'rewards/margins': (rewards_chosen - rewards_rejected).mean(), + } + if self.args.center_rewards_coefficient is not None: + metric['center_rewards_loss'] = center_rewards_loss.detach() + metric = self._all_reduce_metric(metric) + return loss, metric + + def forward_step(self, data_iterator, model): + vp_stage = model.module.module.vp_stage + data = self.get_batch(data_iterator, vp_stage) + data.pop('loss_scale', None) + output_tensor = model(**data) + return output_tensor, partial(self.loss_func, data=data) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/rlhf_mixin.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/rlhf_mixin.py new file mode 100644 index 0000000000000000000000000000000000000000..5cbd83a36f250138befe3f927a374c3d2e228541 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/rlhf_mixin.py @@ -0,0 +1,104 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from contextlib import contextmanager +from megatron.core import mpu +from torch.distributed.nn import all_reduce +from transformers.utils import ContextManagers + +from swift.megatron.model import get_mcore_model +from swift.megatron.utils import load_mcore_checkpoint +from swift.rlhf_trainers.utils import identity_data_collator +from swift.utils import get_current_device, get_logger, safe_snapshot_download +from .base import BaseMegatronTrainer +from .utils import compute_per_token_logps_fn, reconstruct_tensor_cp + +logger = get_logger() + + +class MegatronRLHFTrainer(BaseMegatronTrainer): + + def _load_checkpoint(self): + args = self.args + if args.mcore_ref_model is not None: + load_mcore_checkpoint(args, self.ref_models, load_arg='mcore_ref_model') + if args.mcore_ref_adapter is not None: + load_mcore_checkpoint(args, self.wrapped_models, load_arg='mcore_ref_adapter') + super()._load_checkpoint() + + def prepare_model(self): + super().prepare_model() + args = self.args + self.ref_models = [] + if args.tuner_type == 'full' and args.rlhf_type not in ['rm', 'gkd']: + self.ref_models = get_mcore_model(args, self.template.config) + for ref_model in self.ref_models: + if not args.use_cpu_initialization: + ref_model.to(get_current_device()) + ref_model.requires_grad_(False) + ref_model.eval() + if self.ref_models and args.mcore_ref_model is None: + ref_model_id_or_path = args.ref_model or args.model + ref_model_dir = safe_snapshot_download(ref_model_id_or_path, use_hf=args.use_hf, hub_token=args.hub_token) + self.bridge.load_weights(self.ref_models, ref_model_dir) + if args.tuner_type == 'lora' and args.ref_adapters and args.mcore_ref_adapter is None: + assert len(args.ref_adapters) == 1, 'Currently only support one adapter.' + self.bridge.load_weights( + self.ref_models, args.ref_adapters[0], peft_format=True, adapter_name='ref_adapter') + + def _get_data_collator(self): + if self.args.rlhf_type in ('grpo', 'gkd'): + return identity_data_collator + return super()._get_data_collator() + + @contextmanager + def null_ref_context(self): + args = self.args + contexts = [] + has_ref_adapter = bool(args.mcore_ref_adapter or args.ref_adapters) + if args.tuner_type == 'full': + ref_models = self.ref_models + else: + if not has_ref_adapter: + for m in self.peft_models: + contexts.append(m.disable_adapter()) + ref_models = self.unwrapped_models + with ContextManagers(contexts): + if has_ref_adapter: + for m in self.peft_models: + m.set_adapter('ref_adapter') + yield ref_models + if has_ref_adapter: + for m in self.peft_models: + m.set_adapter('default') + + def get_logps(self, output_tensor, labels, packed_seq_params, per_token=False): + args = self.args + per_token_logps = -output_tensor + loss_mask = labels != -100 + per_token_logps = per_token_logps * loss_mask + num_samples = packed_seq_params.seq_lens.shape[0] if packed_seq_params is not None else labels.shape[0] + if per_token: + if args.context_parallel_size > 1: + per_token_logps = reconstruct_tensor_cp(args.context_parallel_size, per_token_logps, packed_seq_params, + num_samples) + return per_token_logps + + if args.padding_free: + cu_seqlens = packed_seq_params.cu_seqlens_q[:num_samples + 1] // args.context_parallel_size + all_logps = per_token_logps.new_zeros((num_samples, )) + for i in range(num_samples): + start, end = cu_seqlens[i], cu_seqlens[i + 1] + all_logps[i] = per_token_logps[:, start:end].sum() + else: + all_logps = per_token_logps.sum(-1) + if args.context_parallel_size > 1: + all_logps = all_reduce(all_logps, group=mpu.get_context_parallel_group()) + return all_logps + + def compute_per_token_logps(self, model, data_iterator, no_grad=True, temperature=1.0): + return compute_per_token_logps_fn( + model, + self.args, + data_iterator, + temperature=temperature, + no_grad=no_grad, + enable_routing_replay=self.enable_routing_replay) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/rollout_mixin.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/rollout_mixin.py new file mode 100644 index 0000000000000000000000000000000000000000..c9725cbb893bb00a240515c280c703996b6df33a --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/rollout_mixin.py @@ -0,0 +1,1013 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +""" +Megatron Rollout Mixin - Provides vLLM integration for on-policy generation. + +This mixin extracts common vLLM rollout functionality from MegatronGRPOTrainer +to be reused by GKD and other trainers that need online generation. +""" +import base64 +import inspect +import json +import os +import re +import time +import torch +import uuid +from accelerate.utils import broadcast_object_list +from collections import OrderedDict, defaultdict, deque +from contextlib import contextmanager, nullcontext +from copy import copy +from dacite import from_dict +from dataclasses import asdict +from megatron.core import mpu +from megatron.core.rerun_state_machine import RerunDataIterator +from transformers import AutoConfig +from transformers.utils import is_torch_npu_available +from typing import Any, Dict, List, Optional, Tuple, Union + +from swift.infer_engine.protocol import RequestConfig, RolloutInferRequest, RolloutOutput +from swift.megatron.model import get_mcore_model +from swift.rl_core.data import OnPolicySample +from swift.rlhf_trainers.base_rollout_mixin import BaseRolloutTrainerMixin +from swift.rlhf_trainers.gkd_helpers import resolve_dynamic_opd_self_distillation +from swift.rlhf_trainers.utils import (VLLM_LORA_INT_ID, VLLM_LORA_NAME, VLLM_LORA_PATH, FlattenedTensorBucket, + TensorLoRARequest, add_base_layer_suffix_by_param_names, aggressive_empty_cache, + check_vllm_version_ge, expand_vllm_param_name_aliases, finish_vllm_weight_reload, + parse_prompt_logprobs, patch_vllm_load_adapter, + patch_vllm_moe_model_weight_loader, profiling_context, profiling_decorator, + set_expandable_segments, vllm_supports_lora_load_inplace) +from swift.rlhf_trainers.vllm_client import VLLMInferClient +from swift.rollout import MultiTurnScheduler, invoke_async_hook, multi_turns, run_multi_turn +from swift.utils import (JsonlWriter, get_current_device, get_logger, is_last_rank, is_vllm_available, remove_response, + synchronize, to_device) +from .utils import (gather_object, load_megatron_model_to_gpu, load_megatron_optimizer, offload_megatron_model_to_cpu, + offload_megatron_optimizer) + +DataType = List[Dict[str, Union[torch.Tensor, Any]]] +logger = get_logger() + + +def create_rollout_group(trainer) -> torch.distributed.ProcessGroup: + """ + Get or create the rollout process group (TP×PP×CP). + + This is a shared function used by both MegatronRolloutMixin and MegatronGRPOTrainer. + + The rollout group is used for: + 1. Data slicing: distributing rollout data across ranks with same data samples + 2. Gather operations: collecting results from ranks with same data samples + + Note: Groups are created per data parallel index, containing TP×PP×CP ranks each. + This follows Megatron's data_iterator logic where same data_parallel_rank processes + identical data samples. + + Key insight: ranks with the SAME data parallel index process the SAME data samples + and must coordinate for rollout data distribution. + Megatron rank order: TP → CP → EP → DP → PP + + Args: + trainer: Trainer instance with _rollout_group and _rollout_groups_created attributes + + Returns: + The rollout process group for this rank + """ + if trainer._rollout_group is not None: + return trainer._rollout_group + + cp_size = mpu.get_context_parallel_world_size() + if cp_size == 1: + # No CP, use the standard MODEL_PARALLEL_GROUP + trainer._rollout_group = mpu.get_model_parallel_group() + return trainer._rollout_group + + # Use RankGenerator to create rollout groups following Megatron-LM logic + global_rank = torch.distributed.get_rank() + + # Get parallel dimensions + tp_size = mpu.get_tensor_model_parallel_world_size() + pp_size = mpu.get_pipeline_model_parallel_world_size() + dp_size = mpu.get_data_parallel_world_size() + + # Create RankGenerator following Megatron-LM pattern + # Order: tp-cp-ep-dp-pp (default in Megatron-LM) + decoder_rank_generator = mpu.RankGenerator( + tp=tp_size, + ep=1, + dp=dp_size, + pp=pp_size, + cp=cp_size, + order='tp-cp-ep-dp-pp', + rank_offset=0, + ) + + # Create rollout groups based on data consistency from data_iterator + # Same data_parallel_rank processes same data - group ranks with same DP index + if not trainer._rollout_groups_created: + # Use 'tp-cp-ep-pp' to get groups with same DP index (DP is excluded from variation) + dp_groups = decoder_rank_generator.get_ranks('tp-cp-ep-pp') + for dp_group_ranks in dp_groups: + # Sort for consistency + dp_group_ranks = sorted(dp_group_ranks) + group = torch.distributed.new_group(ranks=dp_group_ranks, group_desc='ROLLOUT_GROUP') + + if global_rank in dp_group_ranks: + trainer._rollout_group = group + trainer._rollout_groups_created = True + + return trainer._rollout_group + + +class MegatronRolloutMixin(BaseRolloutTrainerMixin): + + # Per-sample container class; subclasses override (GRPOSample / GKDSample). + sample_cls = OnPolicySample + + def _init_rollout_params(self): + """Initialize rollout generation parameters.""" + args = self.args + # distributed params + self.world_size = torch.distributed.get_world_size() + self.process_index = torch.distributed.get_rank() + self.is_main_process = is_last_rank() + self.device = get_current_device() + + # sampling params + self.temperature = getattr(args, 'temperature', 1.0) + self.max_completion_length = args.max_completion_length + structured_outputs_regex = getattr(args, 'structured_outputs_regex', None) + + self.request_config = RequestConfig( + n=1, + max_tokens=args.max_completion_length, + temperature=args.temperature, + top_p=getattr(args, 'top_p', 1.0), + top_k=getattr(args, 'top_k', -1), + min_p=getattr(args, 'min_p', 0.0), + repetition_penalty=getattr(args, 'repetition_penalty', 1.0), + stop=getattr(args, 'stop_words', None), + return_details=True, + logprobs=True, + structured_outputs_regex=structured_outputs_regex) + + self._last_loaded_step = -1 + self._step = 0 + self._rollout_group = None # Lazily initialized rollout group (TP×PP×CP) + self._rollout_groups_created = False # Flag for group creation (all ranks must create together) + self._bridge = None + + def _replace_data_iterator(self, data_iterator): + if not self.unwrapped_models[0].training: + buffered_inputs = self._build_rollout_buffer(data_iterator) + return RerunDataIterator(iter(buffered_inputs[0])) + + if self._step % self.steps_per_generation == 0: + self._buffered_inputs = self._build_rollout_buffer(data_iterator) + encoded_batches = self._buffered_inputs[self._step % self.steps_per_generation] + self._on_train_step_batch(encoded_batches) + self._step += 1 + return RerunDataIterator(iter(encoded_batches)) + + def _build_rollout_buffer(self, data_iterator) -> List[List[Dict]]: + raise NotImplementedError + + def _on_train_step_batch(self, encoded_batches: List[Dict]) -> None: + pass + + def _get_rollout_group(self): + """Get or create the rollout process group (TP×PP×CP).""" + return create_rollout_group(self) + + def _setup_teacher(self) -> None: + """Resolve teacher mode from args and init the API client when applicable. + + Sets ``teacher_model_server`` / ``use_teacher_api`` / ``_is_self_distillation`` / + ``_teacher_use_disable_adapter`` / ``offload_teacher_model`` / ``_has_teacher``. + Must be called before ``_init_rollout_engine`` (the API client lives on last rank). + """ + args = self.args + self.teacher_model_server = getattr(args, 'teacher_model_server', None) + self.use_teacher_api = self.teacher_model_server is not None + self.offload_teacher_model = args.offload_teacher_model + self._teacher_use_disable_adapter = getattr(args, '_teacher_use_disable_adapter', False) + self._is_self_distillation = (args.teacher_model is None and self.teacher_model_server is None) + self._has_teacher_explicit = ( + args.teacher_model is not None or self.teacher_model_server is not None + or self._teacher_use_disable_adapter) + self._is_dynamic_self_distillation = resolve_dynamic_opd_self_distillation( + has_teacher_explicit=self._has_teacher_explicit, + is_self_distillation=self._is_self_distillation, + ) + self._has_teacher = self._has_teacher_explicit or self._is_dynamic_self_distillation + self.teacher_models = None + self.teacher_configs: list = [] + self.teacher_clients: list = [] + if self.use_teacher_api: + from swift.rlhf_trainers.gkd_helpers import parse_teacher_model_server + self.teacher_configs = parse_teacher_model_server(self.teacher_model_server) + if is_last_rank(): + self.teacher_clients = [VLLMInferClient(base_urls=[cfg.url]) for cfg in self.teacher_configs] + + def _load_teacher_model(self) -> None: + """Load the separate local teacher mcore model (called from ``prepare_model``). + + No-op for the API path, dynamic self-distillation, and same-model LoRA + (disable_adapter) — those reuse the student weights or an external server. + """ + if self.use_teacher_api or self._is_self_distillation or self._teacher_use_disable_adapter: + return + args = self.args + vp_size = getattr(args, 'virtual_pipeline_model_parallel_size', None) + assert vp_size is None or vp_size == 1, 'Teacher distillation does not support VPP.' + self.teacher_hf_config = AutoConfig.from_pretrained(args.teacher_model_dir, trust_remote_code=True) + self.teacher_models = get_mcore_model(args, self.teacher_hf_config) + self.teacher_config = self.teacher_models[0].config + if not args.use_cpu_initialization: + for teacher_model in self.teacher_models: + teacher_model.cuda(torch.cuda.current_device()) + for teacher_model in self.teacher_models: + teacher_model.requires_grad_(False) + teacher_model.eval() + self.teacher_config.bridge.load_weights(self.teacher_models, args.teacher_model_dir) + if self.offload_teacher_model: + self._offload_teacher_models() + logger.info('Teacher models offloaded to CPU to save GPU memory') + + def _offload_teacher_models(self) -> None: + if self.teacher_models and not self.use_teacher_api: + offload_megatron_model_to_cpu(self.teacher_models) + + def _load_teacher_models_to_gpu(self) -> None: + if self.teacher_models and not self.use_teacher_api: + load_megatron_model_to_gpu(self.teacher_models, load_grad=False) + + @contextmanager + def load_teacher_model_context(self): + """Load the teacher to GPU for a forward and offload after (when offloading is on).""" + if not self.offload_teacher_model or self.teacher_models is None: + yield + return + self._load_teacher_models_to_gpu() + try: + yield + finally: + self._offload_teacher_models() + + def _gather_teacher_requests(self, requests: List[RolloutInferRequest]) -> Dict[str, Any]: + """Phase 1 (all ranks, collective): gather this teacher's rollout-group-rank-0 requests. + + Only rollout-group rank 0 contributes (others hold TP/PP/CP replicas of the same data); + contributions are tagged with ``dp_rank`` so segments order by DP rank regardless of world + layout or empty (zero-routed) subsets. Returns a handle with the per-DP-rank segments (for + the main-process infer) plus this DP rank's offset/length (for the later slice). The offset + is the prefix sum of preceding DP ranks' lengths, not ``dp_rank * n``: under multi-teacher + routing each DP rank's subset may differ in length, so equal-length slicing misaligns. + """ + rollout_group = self._get_rollout_group() + rollout_rank = torch.distributed.get_rank(group=rollout_group) + dp_rank = mpu.get_data_parallel_rank() + contribution = (dp_rank, list(requests)) if rollout_rank == 0 else None + + world_size = torch.distributed.get_world_size() + all_contributions = [None] * world_size + torch.distributed.all_gather_object(all_contributions, contribution) + + segments_by_dp = {dp: reqs for c in all_contributions if c is not None for dp, reqs in [c]} + dp_ranks_sorted = sorted(segments_by_dp) + offset = sum(len(segments_by_dp[dp]) for dp in dp_ranks_sorted if dp < dp_rank) + flat_global = [req for dp in dp_ranks_sorted for req in segments_by_dp[dp]] + return {'flat_global': flat_global, 'offset': offset, 'n_local': len(requests)} + + def _infer_teacher_requests(self, handle: Dict[str, Any], topk: int, teacher_client: Optional[Any] = None): + """Phase 2 (main process only, no collective): run the teacher HTTP infer. + + Safe to call concurrently across teachers (distinct clients, no collective inside). + """ + if not handle['flat_global']: # no sample routed to this teacher: skip the empty HTTP call + return [] + client = teacher_client if teacher_client is not None else self.teacher_clients[0] + request_config = RequestConfig(prompt_logprobs=topk, max_tokens=1, temperature=0.0) + responses = client.infer(handle['flat_global'], request_config=request_config, use_tqdm=False) + return [parse_prompt_logprobs(r, topk=topk) for r in responses] + + def _scatter_teacher_parsed(self, handle: Dict[str, Any], parsed_global): + """Phase 3 (all ranks, collective): broadcast the parsed result and slice this rank's part.""" + world_size = torch.distributed.get_world_size() + obj_list = [parsed_global] + torch.distributed.broadcast_object_list(obj_list, src=world_size - 1) + parsed_global = obj_list[0] + offset, n = handle['offset'], handle['n_local'] + return parsed_global[offset:offset + n] + + def _fetch_teacher_parsed_logprobs(self, + requests: List[RolloutInferRequest], + topk: int, + teacher_client: Optional[Any] = None): + """Combined gather→infer→broadcast for a single teacher (serial); returns this rank's slice.""" + handle = self._gather_teacher_requests(requests) + parsed_global = self._infer_teacher_requests(handle, topk, teacher_client) \ + if self.is_main_process else None + return self._scatter_teacher_parsed(handle, parsed_global) + + def _get_local_rollout_batch(self, samples: List[OnPolicySample]) -> List[OnPolicySample]: + """Split batch within rollout group for distributed vLLM generation. + + The batch is evenly split across the rollout group (TP×PP×CP ranks with + the same DP index). This is the base implementation that simply splits + the batch without repetition. + + Subclasses (e.g., GRPO) may override this to implement custom logic like + repeating each prompt num_generations times. + + Note: In Megatron, batch size should always be divisible by rollout group size. + This is ensured by global_batch_size = micro_batch_size * num_microbatches * dp_size, + where rollout_group_size = tp_size * pp_size * cp_size, and world_size = dp_size * rollout_group_size. + + Args: + samples: Full batch of data samples + + Returns: + Local slice of samples for this rank to process + """ + rollout_group = self._get_rollout_group() + rollout_rank = torch.distributed.get_rank(group=rollout_group) + rollout_group_size = torch.distributed.get_world_size(group=rollout_group) + + total_batch_size = len(samples) + assert total_batch_size % rollout_group_size == 0, \ + f'Batch size ({total_batch_size}) must be divisible by rollout group size ({rollout_group_size})' + + per_device_batch_size = total_batch_size // rollout_group_size + start_idx = rollout_rank * per_device_batch_size + end_idx = start_idx + per_device_batch_size + + return samples[start_idx:end_idx] + + def _gather_rollout_results(self, samples: List[OnPolicySample]) -> List[OnPolicySample]: + """Gather rollout results from all ranks in the rollout group. + + Args: + samples: Local rollout results from this rank + + Returns: + Gathered results from all ranks in the rollout group + """ + rollout_group = self._get_rollout_group() + return gather_object(samples, group=rollout_group) + + def _init_rollout_engine(self): + """Initialize vLLM engine for rollout generation.""" + args = self.args + self._init_rollout_params() + + self.vllm_mode = args.vllm_mode + self.vllm_gpu_memory_utilization = args.vllm_gpu_memory_utilization + self.vllm_tensor_parallel_size = args.vllm_tensor_parallel_size + self.use_vllm = args.use_vllm + self.vllm_use_async_engine = False + self.enable_offload = False + self.vllm_version_ge_0_10_2 = check_vllm_version_ge('0.10.2') + self.rollout_enable_lora = False + self.enable_server_multi_turn = False + self.base_sync_done = False + self._cached_vllm_param_names = None + + self._prepare_scheduler() + + if not args.use_vllm: + return + + if args.rlhf_type == 'gkd' and args.lmbda == 0: + return + + if is_torch_npu_available(): + from swift.model.npu_patch.vllm_ascend import validate_vllm_ascend_megatron_lora_training + validate_vllm_ascend_megatron_lora_training(self.unwrapped_models, args) + + if not is_vllm_available(): + raise ImportError('vLLM is not available and `use_vllm` is set to True. ' + 'Please install vLLM with `pip install vllm -U` to use it.') + + if self.vllm_mode == 'server': + # Server mode uses external vLLM server + if self.is_main_process: + self.vllm_client.get_engine_type() + self.vllm_client.reset_mm_cache() + enable_lora = [self.vllm_client.enable_lora] + enable_multi_turn = [self.vllm_client.enable_multi_turn] + else: + enable_lora = [False] + enable_multi_turn = [False] + self.rollout_enable_lora = broadcast_object_list(enable_lora, from_process=self.world_size - 1)[0] + self.enable_server_multi_turn = broadcast_object_list( + enable_multi_turn, from_process=self.world_size - 1)[0] + elif self.vllm_mode == 'colocate': + if self.world_size % self.vllm_tensor_parallel_size != 0: + raise ValueError(f'vllm_tensor_parallel_size ({self.vllm_tensor_parallel_size}) must divide world size ' + f'({self.world_size}) evenly.') + + self.enable_offload = args.offload_model or args.offload_optimizer + context = self.offload_context if self.enable_offload else nullcontext + + with context(): + set_expandable_segments(False) + self.engine = self._prepare_vllm_engine() + self.engine.engine.reset_mm_cache() + if args.sleep_level > 0: + self.engine.engine.sleep(args.sleep_level) + set_expandable_segments(True) + else: + raise ValueError(f'Invalid vllm_mode: {self.vllm_mode}') + + def _prepare_scheduler(self): + """Prepare multi-turn scheduler (shared by GRPO and GKD).""" + args = self.args + + self.multi_turn_scheduler = None + if not hasattr(args, 'multi_turn_scheduler'): + return + + if args.multi_turn_scheduler: + tokenizer = getattr(self, 'processing_class', None) or getattr(self.template, 'tokenizer', None) + if isinstance(args.multi_turn_scheduler, str): + assert args.multi_turn_scheduler in multi_turns + scheduler_kwargs = {'max_turns': args.max_turns, 'tokenizer': tokenizer} + gym_env = getattr(args, 'gym_env', None) + if gym_env is not None: + scheduler_kwargs['gym_env'] = gym_env + multi_turn_scheduler = multi_turns[args.multi_turn_scheduler](**scheduler_kwargs) + self.multi_turn_scheduler: MultiTurnScheduler = multi_turn_scheduler + else: + assert isinstance(args.multi_turn_scheduler, MultiTurnScheduler) + self.multi_turn_scheduler: MultiTurnScheduler = args.multi_turn_scheduler + + def _prepare_vllm_engine(self): + """Create and configure vLLM engine for colocate mode.""" + from vllm.distributed import parallel_state as vllm_ps + + from swift.infer_engine import GRPOVllmEngine + + args = self.args + per_device_batch_size = getattr(args, 'per_device_generation_batch_size', args.micro_batch_size) + max_num_seqs = args.vllm_max_num_seqs or per_device_batch_size * self.vllm_tensor_parallel_size + + vllm_template = copy(self.template) + vllm_template.padding_free = False + vllm_template.sequence_parallel_size = 1 + + logprobs_mode = 'processed_logprobs' if self.vllm_version_ge_0_10_2 else None + + vllm_engine_kwargs = args.vllm_engine_kwargs or {} + load_format = vllm_engine_kwargs.pop('load_format', 'auto') + + if self.args.router_replay_mode == 'R3': + assert check_vllm_version_ge('0.14.0'), \ + 'The enable_return_routed_experts attribute is not supported. Please upgrade vllm to 0.14.0 or higher' + vllm_engine_kwargs['enable_return_routed_experts'] = True + # https://github.com/vllm-project/vllm/pull/39917 + import vllm + from packaging import version + vllm_version = vllm.__version__ + if vllm_version is not None and version.parse('0.21.0rc1') <= version.parse(vllm_version) <= version.parse( + '0.21.0'): + vllm_engine_kwargs.setdefault('async_scheduling', False) + + enable_lora = False + max_loras = 1 + max_lora_rank = args.lora_rank + if args.tuner_type == 'lora' and args.vllm_enable_lora: + enable_lora = True + self.rollout_enable_lora = True + patch_vllm_load_adapter() + logger.info(f'Enabled vLLM LoRA adapter sync with max_lora_rank={args.lora_rank}') + + engine = GRPOVllmEngine( + args.model_info.model_dir, + torch_dtype=args.torch_dtype, + model_type=args.model_type, + use_async_engine=False, + tensor_parallel_size=self.vllm_tensor_parallel_size, + gpu_memory_utilization=self.vllm_gpu_memory_utilization, + enable_prefix_caching=args.vllm_enable_prefix_caching, + max_num_seqs=max_num_seqs, + enforce_eager=args.vllm_enforce_eager, + limit_mm_per_prompt=args.vllm_limit_mm_per_prompt, + enable_sleep_mode=args.sleep_level > 0, + max_model_len=args.vllm_max_model_len, + seed=self.process_index // self.vllm_tensor_parallel_size, + disable_cascade_attn=args.vllm_disable_cascade_attn, + load_format=load_format, + mm_processor_cache_gb=args.vllm_mm_processor_cache_gb, + template=vllm_template, + distributed_executor_backend='external_launcher', + enable_lora=enable_lora, + max_loras=max_loras, + max_lora_rank=max_lora_rank, + engine_kwargs=vllm_engine_kwargs, + logprobs_mode=logprobs_mode) + + if self.vllm_tensor_parallel_size > 1: + self.vllm_tp_group = vllm_ps.get_tp_group().device_group + + return engine + + @profiling_decorator + def _move_model_to_vllm(self): + """Synchronize model weights to vLLM engine. + + - Full sync: when tuner_type != 'lora' (e.g. full, lora_llm), or first sync + (base_sync_done=False), or sleep_level==2, or rollout_enable_lora is disabled. + - Adapter-only sync: when tuner_type == 'lora' with rollout_enable_lora=True and + base weights have already been synced. + """ + args = self.args + tuner_type = args.tuner_type + + if (tuner_type != 'lora' or (not self.base_sync_done or args.sleep_level == 2) or not self.rollout_enable_lora): + self._move_full_model_to_vllm() + else: + self._move_adapter_to_vllm() + + self._reset_vllm_cache() + + def _reset_vllm_cache(self): + # Reset prefix cache and encoder cache + vllm_ge_16 = check_vllm_version_ge('0.16') + if self.vllm_mode == 'server' and self.is_main_process: + self.vllm_client.reset_prefix_cache() + if vllm_ge_16: + self.vllm_client.reset_encoder_cache() + elif self.vllm_mode == 'colocate': + self.engine.engine.reset_prefix_cache() + if vllm_ge_16: + self.engine.engine.reset_encoder_cache() + + def _move_full_model_to_vllm(self): + """Transfer full model weights to vLLM engine. + + For LoRA training (tuner_type == 'lora'): + - When rollout_enable_lora=False: merge LoRA into base, export merged weights, then unmerge. + - When rollout_enable_lora=True: export base weights only (no merge needed), + then follow up with adapter-only sync via _move_adapter_to_vllm. + For lora_llm: always merge LLM LoRA into exported dense weights (vLLM has no separate + adapter pass for this tuner_type; see _move_model_to_vllm). + """ + is_lora_training = self.args.tuner_type in ('lora', 'lora_llm') + is_pure_lora = self.args.tuner_type == 'lora' + should_merge = not self.rollout_enable_lora + if self.args.tuner_type == 'lora_llm' and self.rollout_enable_lora: + logger.warning('lora_llm is not supported with vllm_enable_lora=True. plz set vllm_enable_lora to False') + + try: + if should_merge: + self.merge_lora_adapters() + + self._export_and_load_weights() + + finally: + if should_merge: + self.unmerge_lora_adapters() + + if is_lora_training: + self.base_sync_done = True + if self.rollout_enable_lora and is_pure_lora: + self._move_adapter_to_vllm() + + def _move_adapter_to_vllm(self): + """Transfer only LoRA adapter weights to vLLM engine. + + Uses bridge.export_weights(peft_format=True) to export LoRA delta weights. + Yielded names follow PEFT convention: 'base_model.model..lora_A.weight'. + """ + target_device = 'cpu' if self.args.offload_bridge else None + + with profiling_context(self, 'export_adapter_weights'): + adapter_iterator = self.bridge.export_weights( + self.unwrapped_models, target_device=target_device, peft_format=True) + lora_params = OrderedDict() + for name, tensor in adapter_iterator: + lora_params[name] = tensor.detach() + + peft_config = self.unwrapped_models[0].peft_config.get('default', None) + + if self.vllm_mode == 'colocate': + req_kw = dict( + lora_name=VLLM_LORA_NAME, + lora_int_id=VLLM_LORA_INT_ID, + lora_path=VLLM_LORA_PATH, + peft_config=asdict(peft_config), + lora_tensors=lora_params, + ) + if vllm_supports_lora_load_inplace(): + req_kw['load_inplace'] = True + lora_request = TensorLoRARequest(**req_kw) + self.engine.engine.add_lora(lora_request) + elif self.vllm_mode == 'server' and self.is_main_process: + bucket = FlattenedTensorBucket(named_tensors=list(lora_params.items())) + metadatas = bucket.get_metadata() + flattened_tensor = bucket.get_flattened_tensor() + self.vllm_client.update_adapter_flattened_param(peft_config, metadatas, flattened_tensor) + del bucket, metadatas, flattened_tensor + + del lora_params + + def _export_and_load_weights(self): + """Export weights from Megatron and load to vLLM.""" + target_device = 'cpu' if self.args.offload_bridge else None + + with profiling_context(self, 'export_weights'): + weight_iterator = self.bridge.export_weights(self.unwrapped_models, target_device=target_device) + + if self.rollout_enable_lora: + vllm_param_names = self._get_vllm_param_names_for_mapping() + if vllm_param_names: + weight_iterator = add_base_layer_suffix_by_param_names(weight_iterator, vllm_param_names) + + if self.vllm_mode == 'colocate': + llm_model = self.engine.inner_model + ascend_reload_runner = None + if is_torch_npu_available(): + from swift.model.npu_patch.vllm_ascend import get_vllm_ascend_reload_runner + ascend_reload_runner = get_vllm_ascend_reload_runner(self.engine) + if ascend_reload_runner is not None: + ascend_reload_runner.reload_weights(weights_iterator=weight_iterator, is_checkpoint_format=True) + else: + patch_vllm_moe_model_weight_loader(llm_model) + llm_model.load_weights(weight_iterator) + _model_config = self.engine.engine.model_config + finish_vllm_weight_reload(llm_model, model_config=_model_config, target_device=self.device) + elif self.vllm_mode == 'server': + self._load_weights_to_server_in_buckets(weight_iterator) + if self.is_main_process: + self.vllm_client.process_weights_after_loading() + + def _get_vllm_param_names_for_mapping(self): + """Get vLLM runtime parameter names for base_layer mapping. + + Returns an alias-expanded set so bridge/HF names can match vLLM packed names. + """ + if self.vllm_mode == 'colocate': + llm_model = self.engine.inner_model + raw_names = set(dict[Any, Any](llm_model.named_parameters()).keys()) + return expand_vllm_param_name_aliases(raw_names) + + if self.vllm_mode != 'server' or not self.is_main_process: + return None + + if self._cached_vllm_param_names is None: + keys = self.vllm_client.get_model_state_keys() + self._cached_vllm_param_names = expand_vllm_param_name_aliases(set(keys)) + return self._cached_vllm_param_names + + def _load_weights_to_server_in_buckets(self, weight_iterator): + """Load weights to vLLM server in buckets.""" + bucket_size_mb = int(os.environ.get('SWIFT_UPDATE_WEIGHTS_BUCKET_SIZE', 512)) + bucket_size_bytes = bucket_size_mb * 1024 * 1024 + + current_bucket = [] + current_size = 0 + + for name, param in weight_iterator: + param_size = param.numel() * param.element_size() + current_bucket.append((name, param)) + current_size += param_size + + if current_size > bucket_size_bytes and current_bucket: + self._sync_bucket_to_server(current_bucket) + current_bucket = [] + current_size = 0 + + if current_bucket: + self._sync_bucket_to_server(current_bucket) + + def _sync_bucket_to_server(self, bucket_params: List[Tuple[str, torch.Tensor]]): + """Synchronize a bucket of parameters to vLLM server.""" + if not bucket_params or not self.is_main_process: + return + + # Ensure all async GPU ops (e.g. TP all-gather on NCCL stream from bridge.export_weights) + # are complete before .copy_() reads param data on the default stream. + synchronize() + + bucket = FlattenedTensorBucket(named_tensors=bucket_params) + metadatas = bucket.get_metadata() + flattened_tensor = bucket.get_flattened_tensor() + + self.vllm_client.update_flattened_params(metadatas, flattened_tensor) + + del bucket, metadatas, flattened_tensor + + @profiling_decorator + def _generate_completions(self, samples: List[OnPolicySample]) -> List[OnPolicySample]: + """Generate completions for a batch using vLLM engine. + + Args: + samples: List of OnPolicySample carrying messages + rollout fields + + Returns: + Batch with rollout completion results merged in + """ + samples = self._preprocess_inputs(samples) + + # Wake up engine if sleeping (colocate mode) + if self.vllm_mode == 'colocate' and self.engine.inner_model_executor.is_sleeping: + wake_up_params = inspect.signature(self.engine.engine.wake_up).parameters + kwargs = {'tags': ['weights']} if 'tags' in wake_up_params else {} + aggressive_empty_cache() + self.engine.engine.wake_up(**kwargs) + + # Load model weights if needed + if self._step != self._last_loaded_step or self.args.sleep_level == 2: + self._move_model_to_vllm() + self._last_loaded_step = self._step + + context = self.offload_context if self.enable_offload else nullcontext + with context(): + if (self.vllm_mode == 'colocate' and self.engine.inner_model_executor.is_sleeping + and 'tags' in inspect.signature(self.engine.engine.wake_up).parameters): + aggressive_empty_cache() + set_expandable_segments(False) + self.engine.engine.wake_up(tags=['kv_cache']) + + multi_turn_scheduler = getattr(self, 'multi_turn_scheduler', None) + colocate_multi_turn = ( + multi_turn_scheduler is not None and not getattr(self, 'enable_server_multi_turn', False)) + + if colocate_multi_turn: + requests = self.samples2requests(samples) + invoke_async_hook(multi_turn_scheduler.on_trajectory_start(requests)) + request_config = self._get_request_config() + outputs: List[RolloutOutput] = self._rollout_requests(requests, request_config) + outputs = run_multi_turn( + requests=requests, + first_turn_outputs=outputs, + scheduler=multi_turn_scheduler, + rollout_fn=lambda reqs, cfg: self._rollout_requests(reqs, cfg), + request_config=request_config, + max_turns=self.args.max_turns, + gather_fn=lambda x: gather_object(x, group=self._get_rollout_group()), + ) + else: + # Single-turn rollout (or server multi-turn handled by the engine). + outputs: List[RolloutOutput] = self._rollout(samples) + + # Sleep to release memory + if self.vllm_mode == 'colocate' and self.args.sleep_level > 0: + self.engine.engine.reset_prefix_cache() + self.engine.engine.sleep(level=self.args.sleep_level) + aggressive_empty_cache() + set_expandable_segments(True) + + samples = self._postprocess_rollout_outputs(samples, outputs) + + return samples + + def _rollout_requests(self, requests: List[RolloutInferRequest], + request_config: RequestConfig) -> List[RolloutOutput]: + """Continuation rollout taking already-prepared ``RolloutInferRequest`` objects. + + Used by the multi-turn driver (:func:`swift.rollout.run_multi_turn`) on + every turn after the first. ``_set_inputs_system`` is skipped because the + system prompt is already encoded into ``requests[i].messages``. + """ + if self.vllm_mode == 'server': + return self._server_rollout(requests, request_config) + elif self.vllm_mode == 'colocate': + return self._colocate_rollout(requests, request_config) + raise ValueError(f'Invalid vllm_mode: {self.vllm_mode}') + + def _rollout(self, samples: List[OnPolicySample]) -> List[RolloutOutput]: + """Execute rollout using vLLM engine (system already injected by _preprocess_inputs).""" + request_config = self._get_request_config() + + if self.vllm_mode == 'server': + return self._server_rollout(samples, request_config) + elif self.vllm_mode == 'colocate': + return self._colocate_rollout(samples, request_config) + + def _get_request_config(self) -> RequestConfig: + """Get request config with proper seed for distributed TP groups.""" + request_config = copy(self.request_config) + + if self.vllm_mode == 'colocate' and self.vllm_tensor_parallel_size > 1: + batch_size = getattr(self.args, 'per_device_generation_batch_size', self.args.micro_batch_size) + batch_size *= self.vllm_tensor_parallel_size + request_config.seed = batch_size * (self.process_index // self.vllm_tensor_parallel_size) + + return request_config + + def _server_rollout(self, samples: Union[List[OnPolicySample], List[RolloutInferRequest]], + request_config: RequestConfig) -> List[RolloutOutput]: + """Perform rollout using vLLM server mode.""" + infer_requests = self.samples2requests(samples) + + all_requests = gather_object(infer_requests) + all_requests_lengths = gather_object([len(infer_requests)]) + + if not any(requests for requests in all_requests): + return [] + + if self.is_main_process: + all_outputs: List[RolloutOutput] = self.vllm_client.infer( + infer_requests=all_requests, request_config=request_config) + if len(all_outputs) != len(all_requests): + # Per-turn-split multi-turn (`dynamic_num_samples`) is HF-only + raise NotImplementedError( + 'Per-turn-split multi-turn (dynamic_num_samples) is not supported on Megatron — ' + f'server returned {len(all_outputs)} outputs for {len(all_requests)} requests. ' + 'Return one RolloutOutput per request from MultiTurnScheduler.run (combine turns ' + 'inside response_token_ids: List[List[int]]), or use the HF trainer.') + else: + all_outputs = [None] * len(all_requests) + + all_outputs = broadcast_object_list(all_outputs, from_process=self.world_size - 1) + start_idx = sum(all_requests_lengths[:self.process_index]) + end_idx = start_idx + all_requests_lengths[self.process_index] + outputs = all_outputs[start_idx:end_idx] + + return outputs + + def _colocate_rollout(self, samples: Union[List[OnPolicySample], List[RolloutInferRequest]], + request_config: RequestConfig) -> List[RolloutOutput]: + """Perform co-located rollout with vLLM engine.""" + # Normalize samples (first turn) / RolloutInferRequest (continuation turns) + # into engine-ready requests. + samples = self.samples2requests(samples) + start_idx = 0 + end_idx = len(samples) + + # Handle vLLM tensor parallelism + if self.vllm_tensor_parallel_size > 1: + local_rank_in_group = torch.distributed.get_rank(group=self.vllm_tp_group) + local_input_length = len(samples) + all_input_lengths = [None] * self.vllm_tensor_parallel_size + torch.distributed.all_gather_object(all_input_lengths, local_input_length, group=self.vllm_tp_group) + + start_idx = sum(all_input_lengths[:local_rank_in_group]) + end_idx = start_idx + all_input_lengths[local_rank_in_group] + + gathered = [None for _ in range(self.vllm_tensor_parallel_size)] + torch.distributed.all_gather_object(gathered, samples, group=self.vllm_tp_group) + samples = [p for sublist in gathered for p in sublist] + + outputs: List[RolloutOutput] = self.engine.infer( + infer_requests=samples, request_config=request_config, use_tqdm=False) + + if self.vllm_tensor_parallel_size > 1: + # R3 router replay: vLLM's routing capturer host cache only exists on + # TP rank 0, so non-primary TP ranks have routed_experts=None in outputs. + # Broadcast routed_experts from TP primary to all TP ranks. + if getattr(self.args, 'router_replay_mode', None) == 'R3': + routed_experts_list = [output.response.choices[0].routed_experts for output in outputs] + tp_primary_global_rank = torch.distributed.get_global_rank(self.vllm_tp_group, 0) + torch.distributed.broadcast_object_list( + routed_experts_list, src=tp_primary_global_rank, group=self.vllm_tp_group) + if local_rank_in_group != 0: + for output, experts in zip(outputs, routed_experts_list): + output.response.choices[0].routed_experts = experts + outputs = outputs[start_idx:end_idx] + + return outputs + + def _preprocess_inputs(self, samples: List[OnPolicySample]) -> List[OnPolicySample]: + """Preprocess samples before rollout inference. + + Unified pre-processing (mirrors HF RolloutTrainerMixin._preprocess_inputs): + 1. Insert default system message if absent + 2. Assign unique request_id (for rollout tracking) + 3. Strip any prior assistant response (prompt-only for generation) + """ + samples = self._set_inputs_system(samples) + for s in samples: + if not s.request_id: + s.request_id = f'chatcmpl-{str(uuid.uuid4().hex)}' + remove_response(s.messages) + return samples + + def samples2requests(self, samples: Union[List[OnPolicySample], + List[RolloutInferRequest]]) -> List[RolloutInferRequest]: + """Convert samples into RolloutInferRequest objects. + + Already-built ``RolloutInferRequest`` (multi-turn continuation) pass + through unchanged. The per-sample mapping lives in + ``OnPolicySample.to_infer_request``. + """ + if not samples: + return [] + requests_list = [] + for data in samples: + if isinstance(data, RolloutInferRequest): + requests_list.append(data) + else: + include_extra = bool(getattr(self.args, 'vllm_server_pass_dataset', False)) or bool( + getattr(self, 'multi_turn_scheduler', None)) + requests_list.append(data.to_infer_request(include_extra=include_extra)) + return requests_list + + def _log_completions_from_samples(self, samples: List[OnPolicySample]) -> None: + """Log prompts/completions from sample messages (post-filtering). + + Unlike ``_log_rollout`` (which needs raw RolloutOutput objects), this + method reads completions directly from ``sample.messages[-1]['content']``, + so it works after ``_dynamic_sampling`` has filtered / resampled samples. + """ + if not self.log_completions: + return + if not self.unwrapped_models[0].training: + return + messages = gather_object([s.messages for s in samples]) + completions = [] + for s in samples: + content = s.messages[-1]['content'] if s.messages else '' + if isinstance(content, str): + completions.append(content) + elif isinstance(content, list): + completions.append(self.template.safe_decode(content)) + elif isinstance(content, dict) and 'input_ids' in content: + completions.append(self.template.safe_decode(content['input_ids'])) + else: + completions.append(str(content)) + completions = gather_object(completions) + self._logs['prompt'].extend(self._apply_chat_template_to_messages_list(messages)) + self._logs['completion'].extend(completions) + + def _prepare_logging(self): + """Initialize logging infrastructure (shared by GRPO and GKD).""" + args = self.args + self.log_completions = getattr(args, 'log_completions', False) + self.wandb_log_unique_prompts = getattr(args, 'wandb_log_unique_prompts', False) + self.jsonl_writer = JsonlWriter(os.path.join(args.output_dir, 'completions.jsonl'), write_on_rank='last') + self._last_logged_step = -1 + self._logs = { + 'prompt': deque(), + 'completion': deque(), + } + + def _flush_log_completions(self): + """Flush accumulated completion logs to jsonl/wandb/swanlab.""" + if not (self.log_completions and self.is_main_process and len(self._logs['prompt']) > 0): + return + if self._step == self._last_logged_step: + return + self._last_logged_step = self._step + + table = self._build_log_table() + self.jsonl_writer.append(table) + for val in self._logs.values(): + if isinstance(val, deque): + val.clear() + elif isinstance(val, defaultdict): + for d in val.values(): + d.clear() + + args = self.args + if 'wandb' in args.report_to: + import pandas as pd + import wandb + df = pd.DataFrame(table) + if self.wandb_log_unique_prompts: + df = df.drop_duplicates(subset=['prompt']) + wandb.log({'completions': wandb.Table(dataframe=df)}, commit=False) + if 'swanlab' in args.report_to: + import swanlab + headers = list(table.keys()) + rows = [[table[h][i] for h in headers] for i in range(len(table.get('gen_step', table['prompt'])))] + swanlab.log({'completions': swanlab.echarts.Table().add(headers, rows)}) + + def _build_log_table(self) -> Dict[str, list]: + """Build the completion log table. Subclasses extend with extra columns (rewards/advantages).""" + return { + 'gen_step': [self._step - 1] * len(self._logs['prompt']), + 'prompt': list(self._logs['prompt']), + 'completion': list(self._logs['completion']), + } + + def _apply_chat_template_to_messages_list(self, messages_list): + """Convert messages list to prompt text using template.""" + from swift.template import TemplateInputs + prompts_text = [] + for messages in messages_list: + remove_response(messages) + template_inputs = TemplateInputs.from_dict({'messages': messages}) + res = self.template.encode(template_inputs) + prompts_text.append(self.template.safe_decode(res['input_ids'])) + return prompts_text + + @contextmanager + def offload_context(self): + """Context manager for model/optimizer offloading during vLLM generation.""" + if self.args.offload_model: + offload_megatron_model_to_cpu(self.wrapped_models) + if hasattr(self, 'ref_models') and self.ref_models: + offload_megatron_model_to_cpu(self.ref_models) + if getattr(self, 'optimizer', None) and self.args.offload_optimizer: + offload_megatron_optimizer(self.optimizer) + + try: + yield + finally: + if self.args.offload_model: + load_megatron_model_to_gpu(self.wrapped_models) + if hasattr(self, 'ref_models') and self.ref_models: + load_megatron_model_to_gpu(self.ref_models) + if getattr(self, 'optimizer', None) and self.args.offload_optimizer: + load_megatron_optimizer(self.optimizer) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/trainer.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..a5fa72bc68b186a72fbab86918a41f0615883e6d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/trainer.py @@ -0,0 +1,138 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +import torch.distributed as dist +import torch.nn +from collections import defaultdict +from functools import partial +from megatron.core import mpu +from torch.distributed.nn import all_reduce +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss +from typing import List, Optional + +from swift.utils import get_logger +from .base import BaseMegatronTrainer + +logger = get_logger() + + +class MegatronTrainer(BaseMegatronTrainer): + + def seq_cls_loss_func(self, output_tensor, *, labels: torch.Tensor, packed_seq_params=None, attention_mask=None): + args = self.args + logits = self.get_last_tokens(output_tensor, packed_seq_params, attention_mask) + num_labels = args.num_labels + acc = None + if args.problem_type == 'regression': + loss_fct = MSELoss() + if num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif args.problem_type == 'single_label_classification': + loss_fct = CrossEntropyLoss() + logits = logits.view(-1, num_labels) + labels = labels.view(-1) + loss = loss_fct(logits, labels) + acc = (logits.detach().argmax(dim=-1) == labels).float().mean() + elif args.problem_type == 'multi_label_classification': + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + preds = logits.sigmoid() > 0.5 + acc = (labels == preds).all(dim=-1).float().mean() + metric = {'loss': loss.detach().clone()} + if acc is not None: + metric['acc'] = acc + metric = self._all_reduce_metric(metric) + return loss, metric + + def loss_func(self, + output_tensor: torch.Tensor, + *, + labels: torch.Tensor, + loss_scale: Optional[torch.Tensor] = None, + channels: Optional[List[str]] = None, + packed_seq_params=None): + args = self.args + + losses = output_tensor.float() + loss_mask = labels != -100 + if args.enable_dft_loss: + losses = losses * torch.exp(-losses.detach()) + if loss_scale is not None: + losses = losses * loss_scale + loss = torch.cat([torch.sum(losses * loss_mask).view(1), loss_mask.sum().view(1)]) + + # Reduce loss for logging. + reporting_loss = loss.detach().clone() + torch.distributed.all_reduce(reporting_loss, group=mpu.get_data_parallel_group(with_context_parallel=True)) + + lm_loss = loss[0] + lm_loss = lm_loss.clone() + local_num_tokens = loss[1].detach().clone().to(torch.int) + metrics = {'loss': reporting_loss} + if args.enable_channel_loss: + metrics.update(self._compute_channel_loss(losses, loss_mask, channels, packed_seq_params)) + return (lm_loss, local_num_tokens, metrics) + + def _compute_channel_loss(self, losses, loss_mask, channels, packed_seq_params=None): + args = self.args + metrics = defaultdict(lambda: torch.tensor([0.0, 0.0], dtype=torch.float32, device=torch.cuda.current_device())) + if args.padding_free: + num_samples = packed_seq_params.seq_lens.shape[0] + cu_seqlens = packed_seq_params.cu_seqlens_q[:num_samples + 1] // args.context_parallel_size + device = losses.device + total_len = int(cu_seqlens[-1]) + uniq_channels = [None] if channels is None else list(dict.fromkeys(channels)) + ch_index = {c: i for i, c in enumerate(uniq_channels)} + if channels is None: + seg_ch = torch.zeros(num_samples, dtype=torch.long, device=device) + else: + seg_ch = torch.tensor([ch_index[c] for c in channels], dtype=torch.long, device=device) + token_ch = torch.repeat_interleave(seg_ch, cu_seqlens[1:] - cu_seqlens[:-1], output_size=total_len) + mask = loss_mask[0, :total_len].float() + stats = torch.zeros(len(uniq_channels), 2, dtype=torch.float32, device=device) + stats.index_add_(0, token_ch, torch.stack([losses[0, :total_len].detach().float() * mask, mask], dim=-1)) + for c, idx in ch_index.items(): + metrics[f'loss_{c}'] = stats[idx] + else: + for i in range(losses.shape[0]): + channel = None if channels is None else channels[i] + c_loss = losses[i][loss_mask[i]] + metrics[f'loss_{channel}'][0] += c_loss.detach().sum() + metrics[f'loss_{channel}'][1] += c_loss.shape[0] + + # Synchronize keys to avoid getting stuck. + dp_cp_group = mpu.get_data_parallel_group(with_context_parallel=True) + all_keys = [None] * torch.distributed.get_world_size(group=dp_cp_group) + dist.all_gather_object(all_keys, list(metrics.keys()), group=dp_cp_group) + new_metrics = {} + for key in sorted(set().union(*all_keys)): + new_metrics[key] = metrics[key] + new_metrics = self._all_reduce_metric(new_metrics, torch.distributed.ReduceOp.SUM, group=dp_cp_group) + return new_metrics + + def forward_step(self, data_iterator, model): + vp_stage = model.module.module.vp_stage + data = self.get_batch(data_iterator, vp_stage) + loss_scale = data.pop('loss_scale', None) + channels = data.pop('channel', None) + labels = data.get('labels') + if self.args.task_type == 'seq_cls': + data.pop('labels', None) + output_tensor = model(**data) + packed_seq_params = data.get('packed_seq_params') + if self.args.task_type == 'seq_cls': + loss_func = partial( + self.seq_cls_loss_func, + labels=labels, + packed_seq_params=packed_seq_params, + attention_mask=data.get('attention_mask') + if data.get('attention_mask') is not None else data.get('attention_mask_2d')) + else: + loss_func = partial( + self.loss_func, + labels=labels, + loss_scale=loss_scale, + channels=channels, + packed_seq_params=packed_seq_params) + return output_tensor, loss_func diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..66d1e5dfd9cd903bfa9412581a403b0cd242ff10 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/utils.py @@ -0,0 +1,525 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import gc +import torch +from accelerate.utils import gather as hf_gather +from accelerate.utils import gather_object as hf_gather_object +from contextlib import nullcontext +from dataclasses import dataclass +from megatron.core import mpu +from megatron.core.distributed import DistributedDataParallel as DDP +from megatron.core.optimizer import ChainedOptimizer +from typing import Any, Optional + +from swift.dataloader import DataLoaderDispatcher +from swift.megatron.utils import get_batch_on_this_cp_rank, get_packed_seq_params +from swift.utils import empty_cache, get_current_device, get_logger, to_device + +logger = get_logger() + + +def get_batch_on_this_pp_rank(args, data, vp_stage=None): + if args.task_type == 'causal_lm': + data['labels'] = torch.roll(data['labels'], -1, dims=-1) + if 'loss_scale' in data: + data['loss_scale'] = torch.roll(data['loss_scale'], -1, dims=-1) + batch = to_device(data, get_current_device(), non_blocking=True) + if args.pipeline_model_parallel_size == 1: + return batch + is_pp_last_stage = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage) + if not is_pp_last_stage: + batch['labels'] = None + if 'loss_scale' in batch: + batch['loss_scale'] = None + + return batch + + +def gather(tensor, group: Optional[torch.distributed.ProcessGroup] = None): + if group is None: + return hf_gather(tensor) + size = torch.distributed.get_world_size(group=group) + output = [torch.empty_like(tensor) for _ in range(size)] + torch.distributed.all_gather(output, tensor, group=group, async_op=False) + + return torch.cat(output, dim=0) + + +def gather_object(object: Any, group: Optional[torch.distributed.ProcessGroup] = None): + if group is None: + return hf_gather_object(object) + size = torch.distributed.get_world_size(group=group) + output_objects = [None for _ in range(size)] + torch.distributed.all_gather_object(output_objects, object, group=group) + return [x for y in output_objects for x in y] + + +# code borrowed from verl +@torch.no_grad() +def load_megatron_model_to_gpu(models, load_grad=True, load_frozen_params=True): + for model_chunk in models: + if isinstance(model_chunk, DDP): + model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers] + for buffers in model_chunk_all_buffers: + for buffer in buffers: + # sometimes, we don't want to load grad for pure inference + if load_grad and hasattr(buffer, 'grad_data_size'): + current_storage_size = buffer.grad_data.storage().size() + if current_storage_size == 0 or current_storage_size == buffer.grad_data_size: + buffer.grad_data.storage().resize_(buffer.grad_data_size) + buffer.grad_data.zero_() + else: + # Non-standard layers (e.g. GatedDeltaNet) may have grad + # buffers with mismatched storage size; skip resize and + # zero in-place with current storage. + buffer.grad_data.zero_() + + if buffer.param_data.storage().size() == 0: + buffer.param_data.storage().resize_(buffer.param_data_size) + # copy data from cpu to cuda + buffer.param_data.copy_(buffer.param_data.cpu_data, non_blocking=True) + + if load_frozen_params: + device_id = get_current_device() + for param in model_chunk.module.parameters(): + if not param.requires_grad and param.device.type == 'cpu': + param.data = param.data.to(device_id, non_blocking=True) + else: + # we need this for ref module + device_id = get_current_device() + for _, param in model_chunk.named_parameters(): + param.data = param.data.to(device_id, non_blocking=True) + if param.grad is not None: + param.grad = param.grad.to(device_id, non_blocking=True) + gc.collect() + empty_cache() + + +@torch.no_grad() +def offload_megatron_model_to_cpu(models): + """ + In megatron, the model and optimizer storage are: + - bf16 parameter data chunked in model parallel group + - fp32 grad chunked in model parallel group + - fp32 main_parameter chunked in model and dp group + - fp32 optimizer state chunked in model and dp group + + When using LoRA, frozen base model parameters are NOT managed by DDP buffers. + They must be offloaded separately via direct param iteration. + """ + for model_chunk in models: + if isinstance(model_chunk, DDP): + model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers] + for buffers in model_chunk_all_buffers: + for buffer in buffers: + # offload parameters + if buffer.param_data.storage().size() > 0: + existing = getattr(buffer.param_data, 'cpu_data', None) + if existing is None: + buffer.param_data.cpu_data = torch.empty( + buffer.param_data.size(), + dtype=buffer.param_data.dtype, + device='cpu', + pin_memory=True, + ) + buffer.param_data_size = buffer.param_data.storage().size() + else: + assert existing.shape == buffer.param_data.shape, ( + f'cpu_data shape {tuple(existing.shape)} != ' + f'param_data shape {tuple(buffer.param_data.shape)}; ' + 'reallocating would reintroduce the 2x peak.') + assert existing.dtype == buffer.param_data.dtype, ( + f'cpu_data dtype {existing.dtype} != ' + f'param_data dtype {buffer.param_data.dtype}; ' + 'reallocating would reintroduce the 2x peak.') + # Synchronous D2H copy into the preexisting pinned + # buffer; must complete before resize_(0) frees the + # GPU storage. + buffer.param_data.cpu_data.copy_(buffer.param_data.data, non_blocking=False) + buffer.param_data.storage().resize_(0) + + assert buffer.param_data_size == buffer.param_data.cpu_data.storage().size() + + if buffer.grad_data.storage().size() > 0: + # if the grad_data size is already zero, we assume that it is already offloaded + buffer.grad_data_size = buffer.grad_data.storage().size() + buffer.grad_data.storage().resize_(0) + + for param in model_chunk.module.parameters(): + if not param.requires_grad and param.device.type != 'cpu': + param.data = param.data.to('cpu', non_blocking=True) + else: + # we need this for ref module + for _, param in model_chunk.named_parameters(): + param.data = param.data.to('cpu', non_blocking=True) + if param.grad is not None: + param.grad = param.grad.to('cpu', non_blocking=True) + gc.collect() + empty_cache() + + +@torch.no_grad() +def load_megatron_copy_params(optimizers): + """ + Load optimizer parameters back to GPU. Handles ChainedOptimizer. + + Args: + optimizers: Optimizer or ChainedOptimizer instance. + """ + + def _iter_opts(opt): + if isinstance(opt, ChainedOptimizer): + return opt.chained_optimizers + return [opt] + + def load_tensor_to_gpu(tensor): + if tensor is None: + return + device_id = get_current_device() + tensor.data = tensor.data.to(device_id, non_blocking=True) + + def load_group_to_gpu(group): + if group is None: + return + + if isinstance(group, list): + for param_group in group: + if isinstance(param_group, list): + for param in param_group: + load_tensor_to_gpu(param) + else: + load_tensor_to_gpu(param_group) + else: + load_tensor_to_gpu(group) + + # Load all parameter groups to GPU for each underlying optimizer + + for _opt in _iter_opts(optimizers): + if hasattr(_opt, 'shard_fp32_from_float16_groups'): + load_group_to_gpu(_opt.shard_fp32_from_float16_groups) + + +@torch.no_grad() +def offload_megatron_copy_params(optimizers): + """ + Offload optimizer parameters to CPU. Supports both Megatron optimizers + and `ChainedOptimizer`, which wraps a list of underlying optimizers. + + Args: + optimizers: The optimizer or ChainedOptimizer instance. + """ + + def _iter_opts(opt): + if isinstance(opt, ChainedOptimizer): + return opt.chained_optimizers + return [opt] + + def offload_tensor_to_cpu(tensor): + if tensor is None: + return + tensor.data = tensor.data.to('cpu', non_blocking=True) + + def offload_group_to_cpu(group): + if group is None: + return + + if isinstance(group, list): + for param_group in group: + if isinstance(param_group, list): + for param in param_group: + offload_tensor_to_cpu(param) + else: + offload_tensor_to_cpu(param_group) + else: + offload_tensor_to_cpu(group) + + # Offload all parameter groups to CPU for each underlying optimizer + + for _opt in _iter_opts(optimizers): + if hasattr(_opt, 'shard_fp32_from_float16_groups'): + offload_group_to_cpu(_opt.shard_fp32_from_float16_groups) + + +@torch.no_grad() +def load_megatron_optimizer(optimizers): + + def _iter_opts(opt): + if isinstance(opt, ChainedOptimizer): + return opt.chained_optimizers + return [opt] + + for _opt in _iter_opts(optimizers): + load_megatron_copy_params(_opt) + if _opt.optimizer is not None: + # if we are using HybridDeviceOptimizer, we need to only move gpu optimizer state to gpu + if hasattr(_opt.optimizer, '_move_new_state_to_right_device'): + _opt.optimizer._move_new_state_to_right_device() + else: + opt_state_dict_values = _opt.optimizer.state.values() + for v in opt_state_dict_values: + if 'exp_avg' in v: + v['exp_avg'] = v['exp_avg'].to(get_current_device(), non_blocking=True) + if 'exp_avg_sq' in v: + v['exp_avg_sq'] = v['exp_avg_sq'].to(get_current_device(), non_blocking=True) + gc.collect() + empty_cache() + + +@torch.no_grad() +def offload_megatron_optimizer(optimizers): + + def _iter_opts(opt): + if isinstance(opt, ChainedOptimizer): + return opt.chained_optimizers + return [opt] + + for _opt in _iter_opts(optimizers): + offload_megatron_copy_params(_opt) + # worker may hold zero parameter when enabling custom pipeline layout + if _opt.optimizer is not None: + # HybridDeviceOptimizer: offload all sub-optimizer states to CPU + hdo = _opt.optimizer + if all(hasattr(hdo, attr) for attr in ('sub_optimizers', 'inner_param_to_orig_param', 'state')): + for optimizer in hdo.sub_optimizers: + for param, state in optimizer.state.items(): + for k, v in state.items(): + if not isinstance(v, torch.Tensor): + continue + orig_param = hdo.inner_param_to_orig_param.get(param, param) + hdo.state[orig_param][k] = state[k] = v.to('cpu') + else: + opt_state_dict_values = _opt.optimizer.state.values() + for v in opt_state_dict_values: + if 'exp_avg' in v: + v['exp_avg'] = v['exp_avg'].to('cpu', non_blocking=True) + if 'exp_avg_sq' in v: + v['exp_avg_sq'] = v['exp_avg_sq'].to('cpu', non_blocking=True) + gc.collect() + empty_cache() + + +def log_gpu_memory(prefix: str = '', info_once: bool = False): + log_msg = (f'{prefix} GPU memory: {torch.cuda.memory_allocated() / 1024**3:.2f}GB allocated, ' + f'{torch.cuda.memory_reserved() / 1024**3:.2f}GB reserved') + if info_once: + logger.info_once(log_msg, hash_id=prefix) + else: + logger.info(log_msg) + + +@dataclass +class TrainerState: + should_save: bool = False + should_eval: bool = False + should_log: bool = False + + iteration: int = 0 + consumed_train_samples: int = 0 + # compat transformers + max_steps: Optional[int] = None + + best_metric: Optional[float] = None + best_global_step: Optional[int] = None + last_model_checkpoint: Optional[str] = None + best_model_checkpoint: Optional[str] = None + + @property + def global_step(self) -> int: + return self.iteration + + +class MegatronDataLoaderDispatcher(DataLoaderDispatcher): + + @property + def group(self): + return mpu.get_data_parallel_group() + + +def build_streaming_dataloader(args, dataset, collate_fn): + base_dataloader = torch.utils.data.DataLoader( + dataset, + num_workers=args.dataloader_num_workers, + pin_memory=args.dataloader_pin_memory, + collate_fn=collate_fn, + batch_size=args.micro_batch_size, + prefetch_factor=args.dataloader_prefetch_factor if args.dataloader_num_workers > 0 else None, + persistent_workers=args.dataloader_persistent_workers if args.dataloader_num_workers > 0 else False, + ) + return MegatronDataLoaderDispatcher(base_dataloader) + + +_NPU_ATTENTION_MASK_2D_MODEL_TYPES = {'qwen3_5', 'qwen3_5_moe'} + + +def _should_use_npu_generated_attention_mask(args) -> bool: + from transformers.utils import is_torch_npu_available + if not is_torch_npu_available(): + return False + if args.task_type != 'causal_lm' or args.padding_free: + return False + if getattr(args, 'attention_backend', None) == 'local': + return False + return bool(getattr(args, 'use_flash_attn', False)) + + +def _prepare_npu_generated_attention_mask(batch, *, keep_attention_mask_2d: bool) -> None: + if keep_attention_mask_2d: + attention_mask = batch.get('attention_mask') + if 'attention_mask_2d' not in batch and attention_mask is not None: + batch['attention_mask_2d'] = (attention_mask == 0).sum(dim=(1, 2)) > 0 + else: + batch.pop('attention_mask_2d', None) + batch['attention_mask'] = None + + +def prepare_batch(args, data, vp_stage=None): + """Prepare a micro-batch for Megatron forward: PP slicing, packed_seq_params, CP slicing. + + Extracted from BaseMegatronTrainer._prepare_batch for reuse in ray workers. + """ + batch = get_batch_on_this_pp_rank(args, data, vp_stage=vp_stage) + seq_lens = batch.pop('seq_lens', None) + # Consider compatibility and security. + num_samples = batch.pop('num_samples', None) + if seq_lens is not None: + if num_samples is not None: + assert num_samples == len(seq_lens), ( + f"'num_samples' ({num_samples}) is inconsistent with len(seq_lens) ({len(seq_lens)}).") + num_samples = len(seq_lens) + text_position_ids = batch.pop('text_position_ids', None) + if text_position_ids is None: + text_position_ids = batch.get('position_ids') + if _should_use_npu_generated_attention_mask(args): + _prepare_npu_generated_attention_mask( + batch, keep_attention_mask_2d=getattr(args, 'model_type', None) in _NPU_ATTENTION_MASK_2D_MODEL_TYPES) + else: + batch.pop('attention_mask_2d', None) + if args.padding_free and text_position_ids is not None: + batch['packed_seq_params'] = get_packed_seq_params(args, text_position_ids) + if seq_lens is not None: + batch['packed_seq_params'].seq_lens = torch.tensor(seq_lens, device=text_position_ids.device) + if num_samples is not None: + batch['packed_seq_params'].num_samples = num_samples + batch.setdefault('attention_mask', None) + batch = get_batch_on_this_cp_rank(args, batch) + return batch + + +def compute_per_token_logps_fn(model, args, data_iterator, temperature=1.0, no_grad=True, enable_routing_replay=False): + """Forward pass → logits → temperature-scaled per-token logps. + + Returns: + (per_token_logps, routing_topk_idx) — either may be None on non-last PP stages. + """ + from swift.megatron.utils import (RouterReplayHelper, forward_step_helper, get_local_topk_idx_for_current_rank, + get_router_replay_data, set_router_replay_data) + from .vocab_parallel_utils import compute_logps_and_entropy_from_logits + + data = prepare_batch(args, next(data_iterator)) + data.pop('loss_scale', None) + labels = data.get('labels') + + routing_topk_idx = None + global_topk_idx = data.pop('routed_experts', None) + if enable_routing_replay and RouterReplayHelper.is_replay_forward_action(model.config): + assert global_topk_idx is not None, 'When router_replay_mode = R3, routed_experts must be in data' + routing_topk_idx = get_local_topk_idx_for_current_rank(global_topk_idx, model.config, + data.get('packed_seq_params')) + set_router_replay_data(routing_topk_idx, model.config) + + data_for_forward = {k: v for k, v in data.items() if k != 'labels'} + context = torch.no_grad() if no_grad else nullcontext() + is_training = model.training + if is_training: + model.eval() + try: + with context: + output_tensor = forward_step_helper(model, data_for_forward) + finally: + if is_training: + model.train() + + if enable_routing_replay and RouterReplayHelper.is_r2_record_action(model.config): + routing_topk_idx = get_router_replay_data(model.config) + + if labels is None or output_tensor is None: + return None, routing_topk_idx + + if temperature != 1.0: + output_tensor.div_(temperature) + per_token_logps, _ = compute_logps_and_entropy_from_logits(output_tensor, labels) + + packed_seq_params = data.get('packed_seq_params') + if packed_seq_params is not None: + num_samples = packed_seq_params.seq_lens.shape[0] + else: + input_ids = data.get('input_ids') + num_samples = input_ids.shape[0] if input_ids is not None else labels.shape[0] + + if args.context_parallel_size > 1: + per_token_logps = reconstruct_tensor_cp(args.context_parallel_size, per_token_logps, packed_seq_params, + num_samples) + return per_token_logps, routing_topk_idx + + +def reconstruct_tensor_cp(cp_size, tensor, packed_seq_params, num_samples): + """In CP mode, all_gather and reconstruct full tensor sequences.""" + cp_rank = mpu.get_context_parallel_rank() + + # All-gather across CP ranks + output_list = [torch.empty_like(tensor) for _ in range(cp_size)] + torch.distributed.all_gather(output_list, tensor.contiguous(), group=mpu.get_context_parallel_group()) + output_list[cp_rank] = tensor + + if packed_seq_params is not None: + cu_seqlens_full = packed_seq_params.cu_seqlens_q + cu_seqlens_cp = cu_seqlens_full // cp_size + + # Calculate total packed length + total_packed_len = cu_seqlens_full[num_samples].item() + output_full = tensor.new_zeros(1, total_packed_len) + + # Reconstruct each sequence + for i in range(num_samples): + start_full = cu_seqlens_full[i].item() + end_full = cu_seqlens_full[i + 1].item() + seq_len = end_full - start_full + + # Length of each chunk after CP split + chunk_len = seq_len // cp_size + half_chunk = chunk_len // 2 + + # Concatenate from each CP rank's output (load-balanced split) + for j in range(cp_size): + o = output_list[j][0] + start_cp = cu_seqlens_cp[i].item() + o0 = o[start_cp:start_cp + half_chunk] + o1 = o[start_cp + half_chunk:start_cp + chunk_len] + + # Place back to full sequence + output_full[0, start_full + j * half_chunk:start_full + (j + 1) * half_chunk] = o0 + output_full[0, end_full - (j + 1) * half_chunk:end_full - j * half_chunk] = o1 + else: + # non-padding_free mode: [batch_size, seq_len/cp_size] -> [batch_size, seq_len] + # Each CP rank has chunks split with load-balanced pattern (2*cp_size chunks) + batch_size = tensor.shape[0] + seq_len_per_cp = tensor.shape[1] + full_seq_len = seq_len_per_cp * cp_size + output_full = tensor.new_zeros(batch_size, full_seq_len) + + # Each CP rank j holds chunks j and (2*cp_size - j - 1) from the original 2*cp_size split + # Reconstruct the full sequence by placing chunks back in correct positions + chunk_len = full_seq_len // (2 * cp_size) + + for j in range(cp_size): + o = output_list[j] + # This rank holds 2 chunks: chunk j and chunk (2*cp_size - j - 1) + half_len = seq_len_per_cp // 2 + o0 = o[:, :half_len] + o1 = o[:, half_len:] + + # Place chunk j at position j * chunk_len + output_full[:, j * chunk_len:(j + 1) * chunk_len] = o0 + reverse_idx = 2 * cp_size - j - 1 + output_full[:, reverse_idx * chunk_len:(reverse_idx + 1) * chunk_len] = o1 + + return output_full diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/vocab_parallel_utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/vocab_parallel_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..df070298013211bf9c71e808babc48ce526a31b1 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/trainers/vocab_parallel_utils.py @@ -0,0 +1,236 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Vocabulary-parallel utilities for Tensor Parallelism. + +This module provides utilities for computing log_softmax, entropy, KL divergence, +and other operations across vocab-parallel sharded tensors in Tensor Parallelism (TP). + +When using TP, the vocabulary dimension is sharded across TP ranks. These utilities +correctly handle the distributed computation by: +1. Finding global max via all_reduce (for numerical stability) +2. Computing sum of exp via all_reduce (for normalization) +3. All-reducing partial sums for final results +""" + +import torch +from megatron.core import mpu +from typing import Optional, Tuple + + +def vocab_parallel_log_softmax(logits: torch.Tensor) -> torch.Tensor: + """Compute log_softmax across vocab-parallel sharded logits. + + When using Tensor Parallelism, vocab is sharded across TP ranks. + This function correctly computes log_softmax by: + 1. Finding global max via all_reduce + 2. Computing sum of exp via all_reduce + 3. Computing log_softmax using the global statistics + + Args: + logits: Logits tensor [..., partition_vocab_size] + + Returns: + log_softmax tensor [..., partition_vocab_size] + """ + tp_size = mpu.get_tensor_model_parallel_world_size() + + if tp_size == 1: + return torch.nn.functional.log_softmax(logits, dim=-1) + + tp_group = mpu.get_tensor_model_parallel_group() + + # Step 1: Find global max for numerical stability + logits_max = logits.max(dim=-1, keepdim=True)[0] + torch.distributed.all_reduce(logits_max, op=torch.distributed.ReduceOp.MAX, group=tp_group) + + # Step 2: Compute exp(logits - max) and sum across all TP ranks + exp_logits = torch.exp(logits - logits_max) + sum_exp = exp_logits.sum(dim=-1, keepdim=True) + torch.distributed.all_reduce(sum_exp, op=torch.distributed.ReduceOp.SUM, group=tp_group) + + # Step 3: Compute log_softmax + log_softmax = logits - logits_max - torch.log(sum_exp) + + return log_softmax + + +def vocab_parallel_entropy(log_probs: torch.Tensor, chunk_size: int = 512) -> torch.Tensor: + """Compute entropy from pre-computed vocab-parallel sharded log probabilities. + + When using Tensor Parallelism, vocab is sharded across TP ranks. + This function correctly computes entropy by: + 1. Computing partial entropy = -sum(exp(log_p) * log_p) on each rank's partition + 2. All-reducing the partial entropies to get the global sum. + + Entropy is computed in chunks to reduce memory usage. + + Args: + log_probs: Pre-computed log probabilities tensor [..., partition_vocab_size] + chunk_size: Number of tokens to process per chunk (default: 512) + + Returns: + Entropy tensor [...] (scalar per position) + """ + tp_group = mpu.get_tensor_model_parallel_group() + tp_size = mpu.get_tensor_model_parallel_world_size() + + # Flatten all but the last dimension for chunked processing + original_shape = log_probs.shape[:-1] + vocab_size = log_probs.shape[-1] + log_probs_flat = log_probs.view(-1, vocab_size) # [total_tokens, partition_vocab_size] + total_tokens = log_probs_flat.shape[0] + + entropies_list = [] + for start_idx in range(0, total_tokens, chunk_size): + end_idx = min(start_idx + chunk_size, total_tokens) + log_probs_chunk = log_probs_flat[start_idx:end_idx] # [chunk_size, partition_vocab_size] + + # Compute partial entropy on this rank's vocab partition + # entropy = -sum(p * log_p) = -sum(exp(log_p) * log_p) + probs = torch.exp(log_probs_chunk) + partial_entropy = -(probs * log_probs_chunk).sum(dim=-1) # [chunk_size] + + # All-reduce to get global entropy if using TP + if tp_size > 1: + torch.distributed.all_reduce(partial_entropy, op=torch.distributed.ReduceOp.SUM, group=tp_group) + + entropies_list.append(partial_entropy) + + # Concatenate all chunks and reshape back + entropies = torch.cat(entropies_list, dim=0) + entropies = entropies.view(original_shape) + + return entropies + + +def vocab_parallel_kl_div(input_log_probs: torch.Tensor, target_log_probs: torch.Tensor) -> torch.Tensor: + """Compute KL divergence for vocab-parallel sharded log probabilities. + + KL(target || input) = sum(target_prob * (target_log_prob - input_log_prob)) + = sum(exp(target_log_prob) * (target_log_prob - input_log_prob)) + + Since both log_probs are sharded across TP, we compute the partial sum + on each rank and then all_reduce to get the global sum. + + Args: + input_log_probs: Input log probabilities [..., partition_vocab_size] + target_log_probs: Target log probabilities [..., partition_vocab_size] + + Returns: + KL divergence per position [...], already reduced across TP + """ + tp_group = mpu.get_tensor_model_parallel_group() + + # Compute partial KL on this rank's vocab partition + target_probs = torch.exp(target_log_probs) + partial_kl = (target_probs * (target_log_probs - input_log_probs)).sum(dim=-1) + + if mpu.get_tensor_model_parallel_world_size() > 1: + tp_group = mpu.get_tensor_model_parallel_group() + torch.distributed.all_reduce(partial_kl, op=torch.distributed.ReduceOp.SUM, group=tp_group) + + return partial_kl + + +def vocab_parallel_gather_logps( + logits: torch.Tensor, + labels: torch.Tensor, + log_probs: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Gather log probabilities for target labels from vocab-parallel logits. + + When using TP, each rank only has a partition of the vocabulary. This function: + 1. Computes log_softmax with vocab-parallel support (if log_probs not provided) + 2. Gathers log probs for target tokens + 3. All-reduces to combine results from all TP ranks + + Args: + logits: Logits tensor [batch, seq, partition_vocab_size] + labels: Token labels [batch, seq], -100 for masked positions + log_probs: Pre-computed log_softmax (optional, to avoid recomputation) + + Returns: + per_token_logps: [batch, seq] log probabilities for target tokens + """ + tp_size = mpu.get_tensor_model_parallel_world_size() + tp_rank = mpu.get_tensor_model_parallel_rank() + + # Compute log_probs if not provided + if log_probs is None: + log_probs = vocab_parallel_log_softmax(logits) + + # Get the local vocab range for this TP rank + partition_vocab_size = log_probs.shape[-1] + if tp_size > 1: + vocab_start = tp_rank * partition_vocab_size + vocab_end = vocab_start + partition_vocab_size + # Check which labels fall within this TP rank's vocab partition + local_labels = labels - vocab_start + # Mask for labels within local vocab range + in_range_mask = (labels >= vocab_start) & (labels < vocab_end) + # Clamp local_labels to valid range for gather + local_labels = local_labels.clamp(min=0, max=partition_vocab_size - 1) + else: + local_labels = labels + in_range_mask = torch.ones_like(labels, dtype=torch.bool) + local_labels = local_labels.clamp(min=0) + + # Gather log probs for target tokens + gathered_logps = torch.gather(log_probs, dim=-1, index=local_labels.unsqueeze(-1)).squeeze(-1) + + # For TP: only the rank that owns the target token has the correct log prob + # Other ranks have incorrect values, so we need to zero them out + if tp_size > 1: + gathered_logps = gathered_logps * in_range_mask.float() + # All-reduce to sum contributions from all ranks + # (only one rank has non-zero value for each token) + torch.distributed.all_reduce( + gathered_logps, op=torch.distributed.ReduceOp.SUM, group=mpu.get_tensor_model_parallel_group()) + + # Apply loss mask (labels == -100 are masked) + loss_mask = labels != -100 + per_token_logps = gathered_logps * loss_mask.float() + + return per_token_logps + + +def compute_logps_and_entropy_from_logits( + logits: torch.Tensor, + labels: torch.Tensor, + compute_entropy: bool = False, + entropy_chunk_size: int = 512, +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Compute per-token log probabilities and optionally entropy from logits. + + This is a unified method that efficiently computes both logps and entropy + in a single pass when both are needed, sharing the log_softmax computation. + + Note: In Megatron, labels are already shifted (via torch.roll in get_batch_on_this_tp_rank), + so logits and labels are already aligned. No additional shift is needed here. + + Temperature scaling should be applied by the caller before invoking this function, + so that this function remains a pure computation without side effects on the input. + + Args: + logits: Logits tensor [batch, seq, partition_vocab_size] or [1, total_tokens, partition_vocab_size]. + Should be pre-scaled by temperature if needed. + labels: Token labels [batch, seq] or [1, total_tokens], -100 for masked positions + compute_entropy: Whether to compute entropy (default: False) + entropy_chunk_size: Chunk size for entropy computation (default: 512) + + Returns: + Tuple of: + - per_token_logps: [batch, seq] or [1, total_tokens] log probabilities for target tokens + - per_token_entropy: Same shape as per_token_logps, or None if compute_entropy=False + """ + # Compute log_softmax (shared for both logps and entropy) + log_probs = vocab_parallel_log_softmax(logits) + + # Gather logps for target tokens + per_token_logps = vocab_parallel_gather_logps(logits, labels, log_probs=log_probs) + + # Compute entropy if requested (reuse log_probs to avoid redundant computation) + per_token_entropy = None + if compute_entropy: + per_token_entropy = vocab_parallel_entropy(log_probs, chunk_size=entropy_chunk_size) + + return per_token_logps, per_token_entropy diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/utils/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..46f625339957ee8964aa4fefc7918a6c127dcaab --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/utils/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +from .convert_utils import test_convert_precision +from .megatron_lm_utils import (disable_forward_pre_hook, enable_forward_pre_hook, get_batch_on_this_cp_rank, + get_optimizer_param_scheduler, init_persistent_async_worker, initialize_megatron, + initialize_tp_communicators, load_mcore_checkpoint, maybe_finalize_async_save, + save_mcore_checkpoint, should_disable_forward_pre_hook, warmup_jit_function, wrap_model) +from .parallel_utils import logical_and_across_model_parallel_group, reduce_max_stat_across_model_parallel_group +from .patcher import patch_merge_fn, patch_torch_dist_shard +from .router_replay_utils import (RouterReplayHelper, apply_router_replay_patch, get_local_topk_idx_for_current_rank, + get_router_replay_data, set_router_replay_data) +from .utils import (forward_step_helper, get_packed_seq_params, get_padding_to, prepare_mcore_model, + reconstruct_tensor_cp) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/megatron/utils/convert_utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/utils/convert_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1b48947e90cb7cab204bdbcdd91da2ece912a281 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/megatron/utils/convert_utils.py @@ -0,0 +1,304 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +import math +import torch +import torch.distributed as dist +import torch.nn as nn +from contextlib import contextmanager, nullcontext +from megatron.core import mpu +from megatron.core.extensions.transformer_engine import TEDotProductAttention +from megatron.core.ssm.mamba_context_parallel import _undo_attention_load_balancing +from megatron.core.tensor_parallel import VocabParallelEmbedding +from megatron.core.tensor_parallel.mappings import (gather_from_sequence_parallel_region, + gather_from_tensor_model_parallel_region) +from typing import Any, Dict + +from swift.utils import HfConfigFactory, get_logger, to_device, to_float_dtype +from .megatron_lm_utils import get_batch_on_this_cp_rank +from .utils import forward_step_helper, get_packed_seq_params, get_padding_to + +logger = get_logger() + + +def _test_params_sum(model): + total_sum = 0 + zero_count = 0 + n_parameter = 0 + for n, p in model.named_parameters(): + n_parameter += 1 + sum_ = p.to(device='cuda', dtype=torch.float32).abs().sum().cpu().item() + if sum_ == 0 and '.lora_B.' not in n: + zero_count += 1 + logger.warning(f'n: {n}, sum: {sum_}') + elif math.isnan(sum_) or math.isinf(sum_) or sum_ > 1e10: + logger.warning(f'n: {n}, sum: {sum_}') + else: + total_sum += sum_ + cond = mpu.get_data_parallel_rank() == 0 + logger.info_if(f'n_parameter: {n_parameter}', cond=cond) + logger.info_if(f'total_sum: {total_sum}', cond=cond) + logger.info_if(f'zero_count: {zero_count}', cond=cond) + + +def _find_modules(model, recurse: bool = True, prefix='', ignore_modules=None): + ignore_modules = ignore_modules or [] + for k in ignore_modules: + if prefix.startswith(k): + return [] + else: + named_children = list(model.named_children()) + + modules = [] + for n, module in named_children: + if module.__class__ is nn.ModuleList: + modules += _find_modules(module, False, prefix=f'{prefix}{n}.', ignore_modules=ignore_modules) + elif recurse: + modules += _find_modules(module, prefix=f'{prefix}{n}.', ignore_modules=ignore_modules) + else: + modules.append(module) + if not named_children: + modules.append(model) + return modules + + +@contextmanager +def _model_cpu_forward_context(modules, + torch_dtype=None, + compute_device=None, + share_embedding: bool = False, + target_device='cpu'): + for module in modules: + try: + origin_torch_dtype = next(module.parameters()).dtype + except StopIteration: + pass + else: + break + embeddings = None + if share_embedding: + embeddings = [module for module in modules if isinstance(module, (nn.Embedding, VocabParallelEmbedding))] + + def _to_cuda_hook(module, args): + if compute_device is not None or torch_dtype is not None: + module.to(device=compute_device, dtype=torch_dtype) + args = to_float_dtype(args, dtype=torch_dtype) + return args + + def _to_cpu_hook(module, args, output): + if share_embedding and module in embeddings or 'rotaryemb' in module.__class__.__name__.lower(): + return + module.to(device=target_device, dtype=origin_torch_dtype) + + hooks = [] + for module in modules: + hooks.append(module.register_forward_pre_hook(_to_cuda_hook)) + hooks.append(module.register_forward_hook(_to_cpu_hook)) + try: + yield + finally: + for hook in hooks: + hook.remove() + + +def get_examples(mm_type: str) -> Dict[str, Any]: + if mm_type == 'image': + data = { + 'messages': [{ + 'role': 'user', + 'content': 'describe the image.' + }, { + 'role': + 'assistant', + 'content': + 'The image depicts a close-up of a kitten with striking features. ' + 'The kitten has a white and gray coat with distinct black stripes, ' + 'particularly noticeable on its face and ears. Its eyes are large ' + 'and expressive, with a captivating blue hue that stands out against ' + "the darker fur around them. The kitten's nose is small and pink, " + 'and it has long, delicate whiskers extending from either side of its mouth. ' + "The background is blurred, drawing attention to the kitten's face and " + 'making it the focal point of the image. The overall impression is ' + 'one of cuteness and charm.' + }], + 'images': ['http://modelscope-open.oss-cn-hangzhou.aliyuncs.com/images/cat.png'] + } + elif mm_type == 'audio': + data = { + 'messages': [{ + 'role': 'user', + 'content': '')[-1].strip() + round_messages.append(assistant_no_think) + else: + round_messages.append(deepcopy(msg)) + + return round_messages + + +class MathTipsScheduler(MultiTurnScheduler): + tips_prompt = 'But wait... It seems I made a mistake,' + + def __init__(self, *args, **kwargs): + from swift.rewards.orm import MathAccuracy + super().__init__(*args, **kwargs) + self.acc_func = kwargs.get('acc_function', MathAccuracy()) + # Cache the tokenized tips_prompt length for loss mask computation + self._tips_token_ids = None + + def _get_tips_token_ids(self, tokenizer) -> List[int]: + """Get tokenized tips_prompt (cached for efficiency).""" + if self._tips_token_ids is None: + # Tokenize without special tokens to get the raw token ids + self._tips_token_ids = tokenizer.encode(self.tips_prompt, add_special_tokens=False) + return self._tips_token_ids + + def check_finished(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', + current_turn: int) -> bool: + last_completion = infer_request.messages[-1]['content'] + # we only give tips once + if self.tips_prompt in last_completion: + return True + solution = infer_request.data_dict['solution'] + + acc = self.acc_func([last_completion], [solution])[0] + if acc == 1: + return True + + return super().check_finished(infer_request, response_choice, current_turn) + + def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', + current_turn: int) -> Dict: + completion = response_choice.message.content + response_token_ids = list(response_choice.token_ids) if response_choice.token_ids else [] + + # Extract logprobs from response_choice before any truncation + rollout_logprobs = self._extract_logprobs_from_choice(response_choice) + + # Truncate completion at or tags + truncate_idx = len(completion) + if '' in completion: + truncate_idx = min(truncate_idx, completion.index('')) + if '' in completion: + truncate_idx = min(truncate_idx, completion.index('')) + + if truncate_idx < len(completion): + # Need to truncate token_ids and logprobs as well + truncated_completion = completion[:truncate_idx] + if response_token_ids and self.tokenizer is not None: + # Find the token index corresponding to the truncation point + # by decoding progressively until we reach or exceed the truncation point + token_truncate_idx = len(response_token_ids) + for i in range(1, len(response_token_ids) + 1): + decoded = self.tokenizer.decode(response_token_ids[:i], skip_special_tokens=False) + if len(decoded) >= truncate_idx: + token_truncate_idx = i + break + response_token_ids = response_token_ids[:token_truncate_idx] + # Truncate logprobs to match + if rollout_logprobs: + rollout_logprobs = rollout_logprobs[:token_truncate_idx] + completion = truncated_completion + + # Add tips_prompt + completion += self.tips_prompt + + # Compute loss_mask for tips tokens + # Note: rollout_logprobs should NOT include tips tokens because: + # 1. Tips tokens have loss_mask=0, so their labels will be -100 + # 2. completion_mask = (labels != -100), so tips tokens won't be in completion_mask + # 3. rollout_logprobs must align with completion_mask, not response_token_ids + if response_token_ids and self.tokenizer is not None: + tips_token_ids = self._get_tips_token_ids(self.tokenizer) + # Loss mask: original tokens = 1, tips tokens = 0 + response_loss_mask = [1] * len(response_token_ids) + [0] * len(tips_token_ids) + # Append tips token ids to response + response_token_ids = response_token_ids + tips_token_ids + # Do NOT extend rollout_logprobs for tips tokens - they are masked out in completion_mask + else: + response_loss_mask = [] + + # Update messages + if infer_request.messages[-1]['role'] == 'assistant': + if not infer_request.messages[-1]['content']: + # Multi-turn continuation: pop the dummy input we add in last turn + infer_request.messages.pop(-1) + infer_request.messages[-1]['content'] = completion + else: + infer_request.messages.append({'role': 'assistant', 'content': completion}) + + result = {'infer_request': infer_request} + if response_token_ids: + result['response_token_ids'] = response_token_ids + result['response_loss_mask'] = response_loss_mask + if rollout_logprobs: + result['rollout_logprobs'] = rollout_logprobs + return result + + +class GYMScheduler(MultiTurnScheduler): + """Gym environment-driven scheduler using universal hooks. + + Implements ``on_trajectory_start`` (env.reset) and ``on_turn_end`` (env.step) + to integrate gym environments into the multi-turn protocol. Works in both + server mode (``run()``) and colocate mode (``run_multi_turn()``). + """ + + def __init__(self, infer_engine: Optional['GRPOVllmEngine'] = None, max_turns: Optional[int] = None, **kwargs): + super().__init__(infer_engine, max_turns, **kwargs) + self.gym_env_name = kwargs.get('gym_env', None) + # Per-trajectory state (keyed by uuid) + self._envs: Dict[str, Env] = {} + self._total_rewards: Dict[str, float] = {} + self._step_rewards: Dict[str, List[float]] = {} + self._pending_obs: Dict[str, Optional[str]] = {} + + async def _close_and_remove(self, uuid: str) -> None: + """Close env for a given uuid and remove all associated state.""" + env = self._envs.pop(uuid, None) + if env is not None: + try: + await env.close() + except Exception: + pass + self._total_rewards.pop(uuid, None) + self._step_rewards.pop(uuid, None) + self._pending_obs.pop(uuid, None) + + # ------------------------------------------------------------------ + # Universal async hooks (called by both run() and run_multi_turn()) + # ------------------------------------------------------------------ + async def on_trajectory_start(self, requests: List['RolloutInferRequest']) -> None: + """Create one env per request and seed messages with initial observation.""" + + async def _init_single(req: 'RolloutInferRequest') -> None: + uuid = req.uuid + if uuid in self._envs: + await self._close_and_remove(uuid) + + env_config = (req.data_dict or {}).get('env_config', {}) if hasattr(req, 'data_dict') else {} + env = self._create_env(env_config) + observation, info, system_message = await env.reset(req) + + messages: Messages = [] + if system_message: + messages.append({'role': 'system', 'content': system_message}) + messages.append({'role': 'user', 'content': observation}) + req.messages = messages + + self._envs[uuid] = env + self._total_rewards[uuid] = 0.0 + self._step_rewards[uuid] = [] + self._pending_obs[uuid] = None + + await asyncio.gather(*[_init_single(req) for req in requests]) + + async def on_turn_end(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', + current_turn: int) -> Dict[str, Any]: + """Advance the gym env, accumulate reward, and return done + rollout_infos.""" + uuid = infer_request.uuid + env = self._envs.get(uuid) + if env is None: + return {'done': True, 'rollout_infos': {}} + + next_obs, reward, done, info = await env.step(deepcopy(infer_request.messages)) + self._total_rewards[uuid] = self._total_rewards.get(uuid, 0.0) + float(reward) + self._step_rewards.setdefault(uuid, []).append(float(reward)) + + self._pending_obs[uuid] = None if done else next_obs + + rollout_infos: Dict[str, Any] = { + 'total_reward': self._total_rewards[uuid], + 'step_rewards': list(self._step_rewards.get(uuid, [])), + 'gym_done': done, + } + if done: + await self._close_and_remove(uuid) + + return {'done': done, 'rollout_infos': rollout_infos} + + # ------------------------------------------------------------------ + # Step hook (injects next observation for the next turn) + # ------------------------------------------------------------------ + def step(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', + current_turn: int) -> Dict[str, Any]: + uuid = infer_request.uuid + next_obs = self._pending_obs.get(uuid) + if next_obs is not None: + infer_request.messages.append({'role': 'user', 'content': next_obs}) + self._pending_obs[uuid] = None + return {'infer_request': infer_request} + + # ------------------------------------------------------------------ + # Env helpers + # ------------------------------------------------------------------ + def _create_env(self, env_config: Dict) -> Env: + env_name = env_config.get('name', self.gym_env_name) + if env_name not in envs: + raise ValueError(f"Environment '{env_name}' not found. Available: {list(envs.keys())}") + return envs[env_name](env_config) + + +class OpenEnvScheduler(GYMScheduler): + """GYMScheduler specialised for OpenEnv environments. + + Unlike GYMScheduler which uses async ``Env`` instances, OpenEnvScheduler + uses :class:`OpenEnvWrapper` whose ``reset()`` / ``step()`` / ``close()`` + are **synchronous** (blocking WebSocket I/O). Subclasses that override + ``on_trajectory_start`` / ``on_turn_end`` should wrap sync wrapper calls + with ``asyncio.to_thread()`` to avoid blocking the event loop. + + Action parsing (LLM text → dict) and observation formatting (dict → str) + are handled by overridable :meth:`parse_action` and :meth:`format_observation` + methods, eliminating the need for ``openenv_*`` command-line parameters. + + All OpenEnv configuration (``base_url``, ``system_message``, ``reset_kwargs`` …) + comes from the dataset's per-row ``env_config``. + """ + + def _create_env(self, env_config: Dict) -> Any: + """Create an :class:`OpenEnvWrapper` (not an ``Env`` subclass).""" + from .openenv_wrapper import OpenEnvWrapper + return OpenEnvWrapper(env_config) + + async def _close_and_remove(self, uuid: str) -> None: + """Close wrapper for a given uuid and remove all associated state. + + Wrapper.close() is synchronous; use ``asyncio.to_thread`` to avoid + blocking the event loop. + """ + import asyncio + wrapper = self._envs.pop(uuid, None) + if wrapper is not None: + try: + await asyncio.to_thread(wrapper.close) + except Exception: + pass + self._total_rewards.pop(uuid, None) + self._step_rewards.pop(uuid, None) + self._pending_obs.pop(uuid, None) + + async def on_trajectory_start(self, requests: List['RolloutInferRequest']) -> None: + """Create one wrapper per request, call ``reset()``, seed messages. + + Uses a semaphore to limit concurrent environment creations (default 4) + to avoid overwhelming the OpenEnv server with simultaneous WebSocket connections. + """ + semaphore = asyncio.Semaphore(getattr(self, 'max_concurrent_envs', 4)) + + async def _init_single(req: 'RolloutInferRequest') -> None: + async with semaphore: + uuid = req.uuid + if uuid in self._envs: + await self._close_and_remove(uuid) + + row_env_config = (req.data_dict or {}).get('env_config', {}) if hasattr(req, 'data_dict') else {} + env_config = {**getattr(self, 'env_config_defaults', {}), **row_env_config} + wrapper = self._create_env(env_config) + + obs, metadata = wrapper.reset() + system_message = env_config.get('system_message', '') + + messages: Messages = [] + if system_message: + messages.append({'role': 'system', 'content': system_message}) + messages.append({'role': 'user', 'content': self.format_observation(obs)}) + req.messages = messages + + self._envs[uuid] = wrapper + self._total_rewards[uuid] = 0.0 + self._step_rewards[uuid] = [] + self._pending_obs[uuid] = None + + await asyncio.gather(*[_init_single(req) for req in requests]) + + async def on_turn_end(self, infer_request: 'RolloutInferRequest', response_choice: 'ChatCompletionResponseChoice', + current_turn: int) -> Dict[str, Any]: + """Parse LLM response, call ``wrapper.step()``, accumulate reward.""" + uuid = infer_request.uuid + wrapper = self._envs.get(uuid) + if wrapper is None: + return {'done': True, 'rollout_infos': {}} + + action_text = response_choice.message.content + action_dict = self.parse_action(action_text) + obs, reward, done, metadata = wrapper.step(action_dict) + + self._total_rewards[uuid] = self._total_rewards.get(uuid, 0.0) + float(reward) + self._step_rewards.setdefault(uuid, []).append(float(reward)) + + next_obs = None if done else self.format_observation(obs) + self._pending_obs[uuid] = next_obs + + rollout_infos: Dict[str, Any] = { + 'total_reward': self._total_rewards[uuid], + 'step_rewards': list(self._step_rewards.get(uuid, [])), + 'gym_done': done, + } + if done: + await self._close_and_remove(uuid) + + return {'done': done, 'rollout_infos': rollout_infos} + + def parse_action(self, text: str) -> Dict[str, Any]: + """Parse LLM response text into an OpenEnv action dict. + + Default: try ``json.loads``, fall back to ``{"message": text}``. + """ + text = text.strip() + # Strip markdown code blocks (e.g. ```json ... ```) + if text.startswith('```'): + lines = text.splitlines() + if len(lines) >= 2 and lines[0].startswith('```') and lines[-1].strip().startswith('```'): + text = '\n'.join(lines[1:-1]).strip() + try: + parsed = json.loads(text) + if isinstance(parsed, dict): + return parsed + return {'message': str(parsed)} + except (json.JSONDecodeError, ValueError): + return {'message': text} + + def format_observation(self, observation: Any) -> str: + """Format OpenEnv observation into a string for the LLM. + + Default: ``json.dumps``. Override in subclasses for environment-specific + formatting (e.g. extract a ``"question"`` field). + """ + try: + return json.dumps(observation, ensure_ascii=False, default=str) + except (TypeError, ValueError): + return str(observation) + + +multi_turns = { + 'math_tip_trick': MathTipsScheduler, + 'gym_scheduler': GYMScheduler, + 'openenv_scheduler': OpenEnvScheduler, + 'thinking_tips_scheduler': ThinkingModelTipsScheduler, +} diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/rollout/openenv_wrapper.py b/pipelines/Tr006_V2/project/ms-swift/swift/rollout/openenv_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..b4da01390bec97572c445bc52fabaa178c682c34 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/rollout/openenv_wrapper.py @@ -0,0 +1,115 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +"""Thin wrapper around OpenEnv's GenericEnvClient with reconnection support. + +This wrapper only manages the WebSocket connection and forwards raw actions +to the server. Action parsing (LLM text → dict) and observation formatting +(dict → LLM string) are handled by the :class:`OpenEnvScheduler` subclass, +not here. + +env_config keys consumed: + base_url (str): OpenEnv server URL (e.g. 'http://localhost:8000'). + reset_kwargs (dict, optional): Extra kwargs passed to client.reset(). +""" +import time +from typing import Any, Dict, Tuple + +from swift.utils import get_logger + +logger = get_logger() + +_MAX_RETRIES = 3 +_RETRY_DELAY = 2.0 + + +class OpenEnvWrapper: + """Thin wrapper around ``GenericEnvClient`` with reconnection support. + + Unlike the previous version, this class does **not** inherit from + ``Env`` and does **not** parse LLM text or format observations. + All such logic lives in :class:`~swift.rollout.multi_turn.OpenEnvScheduler` + and its subclasses. + """ + + def __init__(self, env_config: Dict[str, Any]): + self.base_url = env_config.get('base_url', 'http://localhost:8000') + self.reset_kwargs = env_config.get('reset_kwargs', {}) + self._client = None + + # ------------------------------------------------------------------ + # Connection management + # ------------------------------------------------------------------ + + def _ensure_client(self): + """Lazily create the OpenEnv client (sync wrapper).""" + if self._client is None: + from openenv.core.generic_client import GenericEnvClient + + client = GenericEnvClient(base_url=self.base_url) + sync_client = client.sync() + sync_client.__enter__() + self._client = sync_client + return self._client + + def _reconnect_client(self): + """Close old client (recover from connection loss). + + Only cleans up; the next _ensure_client() call will create a new connection. + """ + if self._client is not None: + try: + self._client.__exit__(None, None, None) + except Exception: + pass + self._client = None + + def _call_with_retry(self, fn_name: str, *args, **kwargs): + """Call a client method with retry on connection failure.""" + last_exc = None + for attempt in range(_MAX_RETRIES): + try: + client = self._ensure_client() + return getattr(client, fn_name)(*args, **kwargs) + except Exception as e: + last_exc = e + logger.warning(f'OpenEnv {fn_name} failed (attempt {attempt + 1}/{_MAX_RETRIES}): {e}') + if attempt < _MAX_RETRIES - 1: + time.sleep(_RETRY_DELAY * (attempt + 1)) + self._reconnect_client() + raise last_exc + + # ------------------------------------------------------------------ + # Public API (synchronous — called from async scheduler via run_in_executor + # or directly since GenericEnvClient.sync() is already blocking) + # ------------------------------------------------------------------ + + def reset(self) -> Tuple[Any, Any]: + """Reset the OpenEnv environment. + + Returns: + (raw_observation, metadata) — unformatted, as returned by the server. + """ + result = self._call_with_retry('reset', **self.reset_kwargs) + return result.observation, getattr(result, 'metadata', None) + + def step(self, action_dict: Dict[str, Any]) -> Tuple[Any, float, bool, Any]: + """Execute one step in the OpenEnv environment. + + Args: + action_dict: Pre-parsed action dict (e.g. ``{'answer': '7'}``). + + Returns: + (raw_observation, reward, done, metadata) — unformatted. + """ + result = self._call_with_retry('step', action_dict) + reward = float(result.reward or 0.0) + done = bool(result.done) + return result.observation, reward, done, getattr(result, 'metadata', None) + + def close(self): + """Close the OpenEnv client connection.""" + if self._client is not None: + try: + self._client.__exit__(None, None, None) + except Exception: + logger.debug('OpenEnv client close failed', exc_info=True) + self._client = None diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d2fefbb6d57d89818097ab21a7e638590ce4b92 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +from .sequence_parallel import SequenceParallel, sequence_parallel +from .utils import (ChunkedCrossEntropyLoss, GatherLoss, GatherTensor, SequenceParallelDispatcher, + SequenceParallelSampler) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/ring_utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/ring_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..5b33b5fb155463fc820bb8ab392b32cbff45f226 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/ring_utils.py @@ -0,0 +1,58 @@ +# Some code borrowed from the awesome work: https://github.com/zhuzilin/ring-flash-attention +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +import torch.distributed as dist +from typing import Optional, Tuple + + +class RingComm: + + def __init__(self, process_group: dist.ProcessGroup): + self._process_group = process_group + self._ops = [] + self.rank = dist.get_rank(self._process_group) + self.world_size = dist.get_world_size(self._process_group) + self._reqs = None + + self.send_rank = (self.rank + 1) % self.world_size + self.recv_rank = (self.rank - 1) % self.world_size + + if process_group is not None: + self.send_rank = dist.get_global_rank(self._process_group, self.send_rank) + self.recv_rank = dist.get_global_rank(self._process_group, self.recv_rank) + + def send_recv(self, to_send: torch.Tensor, recv_tensor: Optional[torch.Tensor] = None) -> torch.Tensor: + if recv_tensor is None: + res = torch.empty_like(to_send) + else: + res = recv_tensor + + send_op = dist.P2POp(dist.isend, to_send, self.send_rank, group=self._process_group) + recv_op = dist.P2POp(dist.irecv, res, self.recv_rank, group=self._process_group) + self._ops.append(send_op) + self._ops.append(recv_op) + return res + + def commit(self): + if self._reqs is not None: + raise RuntimeError('commit called twice') + self._reqs = dist.batch_isend_irecv(self._ops) + + def wait(self): + if self._reqs is None: + raise RuntimeError('wait called before commit') + for req in self._reqs: + req.wait() + self._reqs = None + self._ops = [] + + def send_recv_kv( + self, + k: torch.Tensor, + v: torch.Tensor, + k_buffer: Optional[torch.Tensor] = None, + v_buffer: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + next_k, next_v = self.send_recv(k, k_buffer), self.send_recv(v, v_buffer) + self.commit() + return next_k, next_v diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/sequence_parallel.py b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/sequence_parallel.py new file mode 100644 index 0000000000000000000000000000000000000000..41f36067a697ceb5681fa27099fc286b38c52dd8 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/sequence_parallel.py @@ -0,0 +1,731 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import inspect +import math +import torch +import torch.distributed as dist +from functools import lru_cache, partial +from torch.distributed import init_device_mesh +from transformers import PreTrainedTokenizer +from types import SimpleNamespace +from typing import Optional + +from swift.utils import HfConfigFactory, get_cu_seqlens_from_position_ids, get_device, get_dist_setting +from .ulysses import DistributedAttention +from .zigzag_ring_attn import zigzag_ring_flash_attn_varlen_func + + +@lru_cache(maxsize=None) +def get_signature_parameters(fn): + try: + return inspect.signature(fn).parameters + except (TypeError, ValueError): + return None + + +def has_signature_parameter(fn, parameter: str) -> bool: + parameters = get_signature_parameters(fn) + return parameters is not None and parameter in parameters + + +def call_with_supported_kwargs(fn, *args, **kwargs): + parameters = get_signature_parameters(fn) + if parameters is None: + return fn(*args, **kwargs) + if not any(param.kind == inspect.Parameter.VAR_KEYWORD for param in parameters.values()): + kwargs = {key: value for key, value in kwargs.items() if key in parameters} + return fn(*args, **kwargs) + + +def _call_create_causal_mask(fn, config, input_embeds, attention_mask, cache_position_or_past_key_values, *args, + **kwargs): + if has_signature_parameter(fn, 'cache_position'): + return call_with_supported_kwargs( + fn, + config, + input_embeds, + attention_mask, + cache_position_or_past_key_values, + *args, + **kwargs, + ) + if cache_position_or_past_key_values is None and 'past_key_values' in kwargs: + return call_with_supported_kwargs(fn, config, input_embeds, attention_mask, *args, **kwargs) + return call_with_supported_kwargs( + fn, + config, + input_embeds, + attention_mask, + cache_position_or_past_key_values, + *args, + **kwargs, + ) + + +class SequenceParallel: + + _global_inited: bool = False + + def __init__(self): + self.sp_world_size = None + self.dp_world_size = None + self.rp_world_size = None + self.world_size = None + self.model_dtype = None + self.tokenizer = None + self.device_mesh = None + self.num_heads = None + self.causal_mask_func = None + self.extra_kwargs = {} + + @property + def real_position_ids(self) -> torch.Tensor: + """The real position ids, this is different from the position_ids in mrope""" + return self.extra_kwargs.get('text_position_ids') + + def _prepare_flash_attn(self, base_model: torch.nn.Module): + try: + from transformers import masking_utils + + _origin_flash_attention_mask = masking_utils.flash_attention_mask + + def flash_attention_mask(*args, **kwargs): + if self.world_size == 1: + return _origin_flash_attention_mask(*args, **kwargs) + attention_mask = kwargs.get('attention_mask') + if attention_mask is not None: + if attention_mask.all(): + attention_mask = None + + return attention_mask + + masking_utils.flash_attention_mask = flash_attention_mask + masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['flash_attention_2'] = flash_attention_mask + + def sdpa_mask(*args, **kwargs): + if self.world_size == 1: + return masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'](*args, **kwargs) + if 'cache_position' in kwargs: + device = kwargs['cache_position'].device + else: + # transformers>=5.4.0 + device = kwargs['device'] + cache_position = self.real_position_ids[0] + cache_position = self.pad(cache_position, padding_value=-1, position_ids=self.real_position_ids, dim=0) + cache_position = torch.arange(0, cache_position.shape[0], device=device) + kwargs['kv_length'] = cache_position.shape[0] + if 'cache_position' in kwargs: + kwargs['cache_position'] = cache_position + else: + kwargs['q_length'] = kwargs['kv_length'] + return masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'](*args, **kwargs) + + masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping[ + 'sdpa_origin'] = masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] + masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask + + def create_causal_mask(config, + input_embeds, + attention_mask, + cache_position_or_past_key_values=None, + *args, + **kwargs): + if self.world_size == 1: + return _call_create_causal_mask(masking_utils.origin_create_causal_mask, config, input_embeds, + attention_mask, cache_position_or_past_key_values, *args, **kwargs) + input_embeds = torch.ones( + (input_embeds.shape[0], input_embeds.shape[1] * self.sp_world_size, input_embeds.shape[2]), + dtype=input_embeds.dtype, + device=input_embeds.device) + if has_signature_parameter(masking_utils.origin_create_causal_mask, 'cache_position'): + cache_position_or_past_key_values = torch.arange( + 0, + input_embeds.shape[1], + device=input_embeds.device, + ) + return _call_create_causal_mask(masking_utils.origin_create_causal_mask, config, input_embeds, + attention_mask, cache_position_or_past_key_values, *args, **kwargs) + + masking_utils.origin_create_causal_mask = masking_utils.create_causal_mask + masking_utils.create_causal_mask = create_causal_mask + except ImportError: + pass + + if hasattr(base_model, 'language_model'): + text_model = base_model.language_model + else: + text_model = base_model + + from transformers.modeling_flash_attention_utils import is_flash_attn_available + if is_flash_attn_available(): + # TODO this works for multi-modal models like qwen2.5-vl + # SDPA is not supported, because we need to copy the code to our project, which will bring + # more works for maintaining. + from transformers import modeling_flash_attention_utils + from transformers.modeling_flash_attention_utils import _flash_attention_forward + _distributed_flash_attention = DistributedAttention(_flash_attention_forward, self) + + modeling_flash_attention_utils._flash_attention_forward_origin = _flash_attention_forward + + def flash_attention_forward(query_states: torch.Tensor, key_states: torch.Tensor, + value_states: torch.Tensor, attention_mask: Optional[torch.Tensor], q_len, + *args, **kwargs): + if self.world_size == 1: + return _flash_attention_forward(query_states, key_states, value_states, attention_mask, q_len, + *args, **kwargs) + return _distributed_flash_attention(query_states, key_states, value_states, attention_mask, + q_len * self.sp_world_size, *args, **kwargs) + + modeling_flash_attention_utils._flash_attention_forward = flash_attention_forward + + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + def local_flash_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args, + dist_attn, **kwargs): + if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]: + return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query_states, key_states, + value_states, attention_mask, *args, + **kwargs) + if dist_attn.local_attn is None: + + def _attention(query, key, value, *args, **kwargs): + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + if self.rp_world_size is not None and self.rp_world_size > 1: + position_ids = kwargs['position_ids'] + cu_seqlens = get_cu_seqlens_from_position_ids(position_ids) + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + position_ids = self._split_packed(position_ids, cu_seqlens) + mask = position_ids != -1 + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + # this is important + # the length is correct, but the value is wrong, by mask qkv, we do: + # mask the padded values of q and v to zero + # mask the padded values of k to -1e5 + # to keep the attention correct + query, key, value = self._mask_qkv(query, key, value, mask) + output = zigzag_ring_flash_attn_varlen_func( + query, + key, + value, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + causal=module.is_causal, + dropout_p=kwargs.get('dropout', 0.0), + softmax_scale=kwargs.get('scaling', 0.0), + window_size=kwargs.get('sliding_window') or (-1, -1), + group=self.rp_group) + return output + else: + position_ids = kwargs.get('position_ids') + if self.real_position_ids is not None: + position_ids = self.real_position_ids + # Our -1 padding makes transformers' `_is_packed_sequence` see a single varlen sequence + # whenever the batch has one row -- no matter whether padding_free is on -- yet since + # 5.15.0 it derives the lengths from `position_ids == position_ids.min()` instead of `== 0` + # (huggingface/transformers#47525), reading that -1 as the only sequence start. Supply the + # lengths ourselves, and only for that same single-row case. + if position_ids is not None and query.shape[0] == 1: + position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids) + cu_seqlens = get_cu_seqlens_from_position_ids(position_ids) + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + # `query` is (batch, heads, seq, head_dim) here, i.e. after the all-to-all it spans the + # whole sequence, so its length must match the total the lengths account for. + assert query.shape[2] == cu_seqlens[-1], ( + f'query.shape: {tuple(query.shape)}, cu_seqlens: {cu_seqlens.tolist()}') + kwargs['cu_seq_lens_q'] = cu_seqlens + kwargs['cu_seq_lens_k'] = cu_seqlens + kwargs['max_length_q'] = max_seqlen + kwargs['max_length_k'] = max_seqlen + return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query, key, value, *args, + **kwargs)[0] + + dist_attn.local_attn = _attention + + return dist_attn( + query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask, + *args, **kwargs), None + + def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args, + dist_attn, **kwargs): + # Bypass SP logic when world_size == 1 (SP disabled) or module not in text_model + if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]: + return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query_states, key_states, value_states, + attention_mask, *args, **kwargs) + if dist_attn.local_attn is None: + + def _attention(query, key, value, *args, **kwargs): + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + if self.rp_world_size > 1: + raise NotImplementedError('SDPA does not support Ring attention.') + return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query, key, value, *args, **kwargs)[0] + + dist_attn.local_attn = _attention + return dist_attn( + query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask, + *args, **kwargs), None + + ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'] = ALL_ATTENTION_FUNCTIONS['flash_attention_2'] + ALL_ATTENTION_FUNCTIONS['sdpa_origin'] = ALL_ATTENTION_FUNCTIONS['sdpa'] + ALL_ATTENTION_FUNCTIONS['flash_attention_2'] = partial( + local_flash_attn, dist_attn=DistributedAttention(None, self)) + ALL_ATTENTION_FUNCTIONS['sdpa'] = partial(local_sdpa_attn, dist_attn=DistributedAttention(None, self)) + + def _prepare_forward_hook(self, base_model: torch.nn.Module): + + def pre_forward_split_hook(_self, args, kwargs): + if self.world_size == 1: + return args, kwargs + input_ids = kwargs.get('input_ids', None) + inputs_embeds = kwargs.get('inputs_embeds', None) + position_ids = kwargs['position_ids'] + attention_mask = kwargs.get('attention_mask', None) + if hasattr(_self, 'language_model'): + embed_tokens = getattr(_self.language_model, 'embed_tokens', None) + else: + embed_tokens = getattr(_self, 'embed_tokens', None) + input_ids, inputs_embeds, _, position_ids, attention_mask, _, _ = self.pad_and_split_inputs( + input_ids, + inputs_embeds, + None, + position_ids, + attention_mask, + None, + embed_tokens=embed_tokens, + real_position_ids=self.real_position_ids) + kwargs['input_ids'] = input_ids + kwargs['inputs_embeds'] = inputs_embeds + kwargs['position_ids'] = position_ids + kwargs['attention_mask'] = attention_mask + return args, kwargs + + base_model.register_forward_pre_hook(pre_forward_split_hook, with_kwargs=True) + + def _prepare_moe_aux_loss(self, base_model: torch.nn.Module): + from .utils import GatherLoss + + def moe_aux_loss_hook(module, args, kwargs, output): + router_logits = getattr(output, 'router_logits', None) + if router_logits is None: + return output + + attention_mask = kwargs['attention_mask'] + if attention_mask is None: + batch_size = 1 + else: + batch_size = attention_mask.shape[0] + + assert router_logits[0].shape[0] % batch_size == 0 + seq_len = router_logits[0].shape[0] // batch_size + + _gathered_logits = [] + for i in range(batch_size): + _slice = slice(i * seq_len, (i + 1) * seq_len) + _bs_logits = [logit[_slice] for logit in router_logits] + compute_device = _bs_logits[0].device + _bs_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in _bs_logits], dim=0) + _bs_logits, _ = GatherLoss.apply(_bs_logits, None, 1, self.real_position_ids) + _gathered_logits.append(_bs_logits) + router_logits = torch.stack(_gathered_logits, dim=0) + if self.real_position_ids is not None: + router_logits = router_logits[:, :, :self.real_position_ids.shape[1], :] + output['router_logits'] = tuple( + [logit.reshape(-1, logit.shape[-1]) for logit in router_logits.split(1, dim=1)]) + return output + + base_model.register_forward_hook(moe_aux_loss_hook, with_kwargs=True) + + def prepare(self, sp_size: int, model: torch.nn.Module, tokenizer: PreTrainedTokenizer, padding_free: bool): + from swift.model import get_llm_model + self.num_heads = HfConfigFactory.get_config_attr(model.config, 'num_key_value_heads') + if self.num_heads is None: + self.num_heads = HfConfigFactory.get_config_attr(model.config, 'num_attention_heads') + assert self.num_heads is not None, 'Cannot find num_heads config in config.json' + self.padding_free = padding_free + self.world_size = sp_size + + llm_model = get_llm_model(model) + + if hasattr(llm_model, 'language_model'): + if hasattr(llm_model.language_model, '_update_causal_mask'): + self.causal_mask_func = llm_model.language_model._update_causal_mask + else: + if hasattr(llm_model, '_update_causal_mask'): + self.causal_mask_func = llm_model._update_causal_mask + + if not SequenceParallel._global_inited: + # these operations are global initializations and patches + self._init_device_mesh() + self._prepare_flash_attn(llm_model) + SequenceParallel._global_inited = True + + self._prepare_forward_hook(llm_model) + + if model.model_info.is_moe_model: + self._prepare_moe_aux_loss(llm_model) + + self.model_dtype = next(model.parameters()).dtype + self.tokenizer = tokenizer + if self.rp_world_size > 1 and not self.padding_free: + raise NotImplementedError( + f'The world_size {self.world_size} needs ulysses/ring-attention, which needs --padding_free true') + + def _mask_qkv(self, query, key, value, mask): + mask = mask.unsqueeze(2).unsqueeze(3) + query = query * mask + value = value * mask + mask = (~mask) * -1e5 # for bf16 + key = key + mask.to(key.dtype) + return query, key, value + + def pad(self, tensor, padding_value, position_ids=None, dim=1): + """Pad tensor for sequence parallel""" + if self.rp_world_size > 1: + world_size = self.world_size * 2 + else: + world_size = self.world_size + + def _do_pad(tensor): + length = tensor.shape[dim] + pad_num = world_size - (length % world_size) + if pad_num == 0 or pad_num == world_size: + return tensor + if not isinstance(padding_value, torch.Tensor): + # ids + pad_shape = ((*tensor.shape[:dim], pad_num, *tensor.shape[dim + 1:]) if dim != -1 else + (*tensor.shape[:dim], pad_num)) + pad = torch.full(pad_shape, padding_value, dtype=tensor.dtype, device=tensor.device) + tensor = torch.cat([tensor, pad], dim=dim) + else: + # For embeddings + tensor = torch.cat([tensor, padding_value.unsqueeze(0).repeat(tensor.shape[0], pad_num, 1)], dim=dim) + return tensor + + if position_ids is not None and self.rp_world_size > 1: + cu_seqlens = get_cu_seqlens_from_position_ids(position_ids) + all_tensors = [] + for i in range(len(cu_seqlens) - 1): + if dim == 1: + sub_tensor = tensor[:, cu_seqlens[i]:cu_seqlens[i + 1]] + elif dim == -1: + sub_tensor = tensor[..., cu_seqlens[i]:cu_seqlens[i + 1]] + else: + raise NotImplementedError() + all_tensors.append(_do_pad(sub_tensor)) + tensor = torch.cat(all_tensors, dim=dim) + + return _do_pad(tensor) + + def gather(self, local_output, dim: int, position_ids=None): + """Gather tensor for sequence parallel - reverse of split""" + if self.world_size == 1: + return local_output + + if self.rp_world_size > 1: + input_dim = local_output.dim() + assert input_dim >= 2 + + if position_ids is not None: + position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids) + + # Step 1: Gather from all sequence parallel ranks + # Each sp_rank has its own piece, we need to gather them first + gathered_sp = [torch.zeros_like(local_output) for _ in range(self.sp_world_size)] + torch.distributed.all_gather(gathered_sp, local_output.contiguous(), group=self.sp_group) + + # Concatenate the sp pieces to form the complete chunk for this rp_rank + rp_chunk = torch.cat(gathered_sp, dim=dim) + + # Step 2: Gather all rp chunks + gathered_rp = [torch.zeros_like(rp_chunk) for _ in range(self.rp_world_size)] + torch.distributed.all_gather(gathered_rp, rp_chunk, group=self.rp_group) + + cu_seqlens = get_cu_seqlens_from_position_ids(position_ids) + all_tensor_length = [] + for i in range(len(cu_seqlens) - 1): + length = cu_seqlens[i + 1] - cu_seqlens[i] + padding_length = math.ceil(length / (self.world_size * 2)) * (self.world_size * 2) + all_tensor_length.append(padding_length) + + full_output = torch.zeros( + [local_output.shape[0], sum(all_tensor_length), *local_output.shape[2:]], device=local_output.device) + for idx_rp, rp_tensor in enumerate(gathered_rp): # rp world size + # re-group the zigzag to the correct order + accumulated_length = 0 + for idx_seq, length in enumerate(all_tensor_length): # sequence number + local_length = length // self.rp_world_size + local_tensor = rp_tensor[:, accumulated_length:local_length + accumulated_length] + chunk_size = local_length // 2 + left_idx = accumulated_length * self.rp_world_size + idx_rp * chunk_size + right_idx = accumulated_length * self.rp_world_size + (idx_rp + 1) * chunk_size + full_output[:, left_idx:right_idx] = local_tensor[:, :chunk_size] + left_idx = accumulated_length * self.rp_world_size + (2 * self.rp_world_size - idx_rp + - 1) * chunk_size + right_idx = accumulated_length * self.rp_world_size + (2 * self.rp_world_size - idx_rp) * chunk_size + full_output[:, left_idx:right_idx] = local_tensor[:, chunk_size:] + accumulated_length += local_length + + return full_output.contiguous() + else: + gathered_sp = torch.empty( + [local_output.shape[0] * self.sp_world_size] + list(local_output.shape[1:]), + dtype=local_output.dtype, + device=local_output.device) + dist.all_gather_into_tensor(gathered_sp, local_output, group=self.sp_group) + gathered_sp = torch.cat(gathered_sp.split(local_output.shape[0], dim=0), dim=dim) + return gathered_sp.contiguous() + + def _split_packed(self, value, cu_seqlens, dim=1): + """Split and re-group in zigzag""" + local_values = [] + for i in range(len(cu_seqlens) - 1): + start, end = cu_seqlens[i], cu_seqlens[i + 1] + if dim == 1: + sub_value = value[:, start:end] + elif dim == -1: + sub_value = value[..., start:end] + else: + raise NotImplementedError() + local_value = sub_value.chunk(2 * self.rp_world_size, dim=dim) + local_values.extend([ + local_value[self.rp_rank], + local_value[2 * self.rp_world_size - 1 - self.rp_rank], + ]) + return torch.cat(local_values, dim=dim).contiguous() + + def split(self, input, dim: int, position_ids=None): + """Split tensor for sequence parallel""" + if self.world_size == 1: + return input + + if self.rp_world_size > 1: + input_dim = input.dim() + assert input_dim >= 2 + cu_seqlens = get_cu_seqlens_from_position_ids(position_ids) + assert torch.all(cu_seqlens % (2 * self.rp_world_size) == 0) + value_chunks = self._split_packed(input, cu_seqlens, dim=dim) + local_value = value_chunks.chunk(self.sp_world_size, dim=dim)[self.sp_rank].contiguous() + return local_value + else: + rank = self.sp_rank + dim_size = input.size(dim) + assert dim_size % self.sp_world_size == 0, ( + f'The dimension to split ({dim_size}) is not a multiple of ' + f'world size ({self.sp_world_size}), cannot split tensor evenly') + + tensor_list = torch.split(input, dim_size // self.sp_world_size, dim=dim) + output = tensor_list[rank].contiguous() + return output + + def pad_and_split_mm_tokens(self, visual_mask, mm_embeds): + input_ids = self.extra_kwargs['input_ids'] + empty_embeds = torch.empty( + (input_ids.shape[0], input_ids.shape[1], mm_embeds.shape[-1])).to(mm_embeds.device).to(mm_embeds.dtype) + empty_embeds[visual_mask] = mm_embeds + + embeds = SimpleNamespace(weight=mm_embeds) + + _, split_input_embeds, _, _, _, _, extra_values = self.pad_and_split_inputs( + None, + empty_embeds, + None, + None, + None, + None, + embeds, + self.real_position_ids, + extra_split_values=[(visual_mask, 0, -1)]) + visual_mask = extra_values[0] + return visual_mask, split_input_embeds[visual_mask] + + def pad_and_split_inputs(self, + input_ids, + input_embeds, + labels, + position_ids, + attention_mask, + loss_scale, + embed_tokens=None, + real_position_ids=None, + extra_split_values=None): + """Common implementation for padding and splitting inputs + + When a sequence comes, it will be split into rp_world_size * 2 sub tensors, and group them as the + zigzag order. So we get rp_world_size tensors, then we split each tensor to sp_world_size ones. + So, we should first pad the original sequence to the length can be divided by 2 * world_size, then re-group it. + + Only support padding_free for ring-attention, because non-padding_free mode needs another pad/split workflow + man, that's a lot of work... + + Args: + input_ids: input_ids + input_embeds: input_embeds + labels: labels + position_ids: position_ids or, position_ids for mrope + attention_mask: attention_mask + loss_scale: loss_scale + embed_tokens: embed_tokens + real_position_ids: the real position_ids to represent the seq length information + extra_split_values: List of Tuples for extra split values, e.g.: (tensor, pad_value, split_dim) + """ + tokenizer = self.tokenizer + real_position_ids = real_position_ids if real_position_ids is not None else position_ids + extra_values = [] + batch_size = input_ids.shape[ + 0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else None + if real_position_ids is not None and batch_size is not None and real_position_ids.shape[0] == batch_size: + # TODO clone everytime, but the position_ids is a small tensor + self.extra_kwargs['text_position_ids'] = real_position_ids.clone() + if input_ids is not None: + input_ids = self.pad(input_ids, padding_value=tokenizer.pad_token_id, position_ids=real_position_ids) + self.extra_kwargs['input_ids'] = input_ids.clone() + if input_embeds is not None: + pad_emb = torch.zeros( + (1, embed_tokens.weight.shape[-1])).to(embed_tokens.weight.device).to(embed_tokens.weight.dtype) + input_embeds = self.pad(input_embeds, padding_value=pad_emb, position_ids=real_position_ids) + batch_size = input_ids.shape[ + 0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else 1 + if position_ids is not None: + position_ids = self.pad(position_ids, padding_value=-1, position_ids=real_position_ids, dim=-1) + if labels is not None: + labels = self.pad(labels, padding_value=-100, position_ids=real_position_ids) + if loss_scale is not None: + loss_scale = self.pad(loss_scale, padding_value=0., position_ids=real_position_ids) + if real_position_ids is not None: + real_position_ids = self.pad(real_position_ids, padding_value=-1, position_ids=real_position_ids) + if (input_ids is not None or input_embeds is not None) and batch_size > 1: + # not padding_free, so not ring-attention + inputs = input_ids if input_ids is not None else input_embeds + attn_shape = inputs.shape[1] # The sequence length + if attention_mask is None: + attention_mask = torch.ones_like(real_position_ids) + # no need position_ids here, because padding_free does not need attention_mask, + # so this is not ring-attention + attention_mask = self.pad(attention_mask, padding_value=0) + cache_position = torch.arange(0, attn_shape, device=inputs.device) + # pad attention mask to 4d to avoid calculation errors + if hasattr(self, 'causal_mask_func') and self.causal_mask_func is not None: + attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype), cache_position, + None, None) + if extra_split_values is not None: + for (tensor, pad_value, split_dim) in extra_split_values: + extra_values.append( + self.pad(tensor, padding_value=pad_value, position_ids=real_position_ids, dim=split_dim)) + if input_ids is not None: + input_ids = self.split(input_ids, dim=1, position_ids=real_position_ids) + if input_embeds is not None: + input_embeds = self.split(input_embeds, dim=1, position_ids=real_position_ids) + if labels is not None: + labels = torch.roll(labels, shifts=-1, dims=-1) + labels = self.split(labels, dim=-1, position_ids=real_position_ids) + if loss_scale is not None: + loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1) + loss_scale = self.split(loss_scale, dim=-1, position_ids=real_position_ids) + + if position_ids is not None: + position_ids = self.split(position_ids, dim=-1, position_ids=real_position_ids) + if extra_split_values is not None: + for i in range(len(extra_values)): + extra_values[i] = self.split( + extra_values[i], dim=extra_split_values[i][2], position_ids=real_position_ids) + return input_ids, input_embeds, labels, position_ids, attention_mask, loss_scale, extra_values + + def _gather_object_dp(self, input_data): + """Gather object for data parallel""" + input_data_list = [None] * self.dp_world_size + dist.all_gather_object(input_data_list, input_data, group=self.dp_group) + return [x for y in input_data_list for x in y] + + def _init_device_mesh(self): + """Initialize device mesh for sequence and ring parallel. + + The logic is unified: + 1. Determine the Sequence Parallel (SP) size first based on GCD to satisfy constraints. + 2. Allocate all remaining model parallelism to Ring Parallel (RP). + """ + _, _, world_size, _ = get_dist_setting() + self.dp_world_size = world_size // self.world_size + # SP size is the GCD of num_heads and world_size, guaranteeing it divides both. + self.sp_world_size = math.gcd(self.num_heads, self.world_size) + # RP takes the remaining factor so all model-parallel GPUs are used. + self.rp_world_size = self.world_size // self.sp_world_size + + if self.rp_world_size > 1: + mesh_shape = (self.dp_world_size, self.rp_world_size, self.sp_world_size) + mesh_dim_names = ('data', 'ring', 'sequence') + else: + mesh_shape = (self.dp_world_size, self.sp_world_size) + mesh_dim_names = ('data', 'sequence') + self.device_mesh = init_device_mesh( + get_device().split(':')[0], mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names) + + def _dim_group(self, dim_name: str): + """Return the process group of the given mesh dim, or None if absent.""" + if not self.device_mesh or dim_name not in self.device_mesh.mesh_dim_names: + return None + return self.device_mesh[dim_name].get_group() + + def _dim_rank(self, dim_name: str, default: int) -> int: + """Return the rank within the given mesh dim, or `default` if absent.""" + group = self._dim_group(dim_name) + return dist.get_rank(group) if group is not None else default + + @property + def sp_group(self): + """Return the sequence parallel group""" + return self._dim_group('sequence') + + def enabled(self): + return self.world_size is not None and self.world_size > 1 + + @property + def sp_rank(self): + """Return the sequence parallel rank""" + return self._dim_rank('sequence', default=0) + + @property + def dp_group(self): + """Return the data parallel group""" + return self._dim_group('data') + + @property + def dp_rank(self): + """Return the data parallel rank""" + return self._dim_rank('data', default=0) + + @property + def rp_group(self): + """Return the ring parallel group""" + return self._dim_group('ring') + + @property + def rp_rank(self): + """Return the ring parallel rank""" + return self._dim_rank('ring', default=-1) + + def prepare_inputs(self, inputs): + """Prepare inputs + + 1. set extra_kwargs['text_position_ids'] + 2. split labels + """ + position_ids = inputs.get('text_position_ids') + input_ids = inputs.get('input_ids') + if position_ids is None: + position_ids = inputs.get('position_ids') + if position_ids is not None and input_ids is not None and position_ids.shape[0] == input_ids.shape[0]: + self.extra_kwargs['text_position_ids'] = position_ids.clone() + if input_ids is not None: + self.extra_kwargs['input_ids'] = input_ids.clone() + if 'labels' in inputs: + labels = inputs['labels'] + _, _, labels, _, _, _, _ = self.pad_and_split_inputs( + None, None, labels, None, None, None, real_position_ids=position_ids) + inputs['labels'] = labels + + +sequence_parallel = SequenceParallel() diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/ulysses.py b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/ulysses.py new file mode 100644 index 0000000000000000000000000000000000000000..5db7e2183ede3ee2cd60780b2e1f37a67725ca16 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/ulysses.py @@ -0,0 +1,147 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +import torch.distributed as dist +from typing import Any, Tuple + + +# Code borrowed from deepspeed, here is why: +# 1. Reduce the dependency +# 2. The original code is complex +def _generate_layout_params(scatter_idx, seq_world_size, input): + if scatter_idx < 2: + bs, global_seq_len, num_local_head, head_dim = input.shape + pre_all2all_inp_shape = [bs, seq_world_size, global_seq_len // seq_world_size, num_local_head, head_dim] + pre_all2all_permute_idx = (1, 0, 2, 3, 4) + + post_all2all_permute_idx = (1, 2, 0, 3, 4) + post_all2all_res_shape = [bs, global_seq_len // seq_world_size, seq_world_size * num_local_head, head_dim] + else: + bs, local_seq_len, num_total_head, head_dim = input.shape + assert num_total_head % seq_world_size == 0, (f'Number of heads ({num_total_head}) must be divisible ' + f'by the sequence parallel size ({seq_world_size})!') + pre_all2all_inp_shape = [bs, local_seq_len, seq_world_size, num_total_head // seq_world_size, head_dim] + pre_all2all_permute_idx = (2, 0, 1, 3, 4) + + post_all2all_permute_idx = (1, 0, 2, 3, 4) + post_all2all_res_shape = [bs, seq_world_size * local_seq_len, num_total_head // seq_world_size, head_dim] + + return pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape + + +def post_all2all(permute_idx, res_shape): + """ + Post-processing function for `all2all` communication. + """ + + def post_func(input): + if permute_idx is not None: + input = input.permute(permute_idx).contiguous() + output = input.reshape(res_shape).contiguous() + + return output + + return post_func + + +def pre_all2all_fun(permute_idx, inp_shape, input): + """ + Pre-processing function for `all2all` communication. + """ + input_t = input.reshape(inp_shape).contiguous() + if permute_idx is not None: + input_t = input_t.permute(permute_idx).contiguous() + return input_t + + +def single_all_to_all(input, scatter_idx, gather_idx, group, **kwargs): + seq_world_size = dist.get_world_size(group) + num_heads = input.shape[2] + if num_heads % seq_world_size != 0 and not scatter_idx < 2: + raise NotImplementedError(f'num_heads {num_heads} cannot be split by sp world size {seq_world_size}') + pre_all2all_permute_idx, pre_all2all_inp_shape, post_all2all_permute_idx, post_all2all_res_shape = ( + _generate_layout_params(scatter_idx, seq_world_size, input)) + + input_t = pre_all2all_fun(pre_all2all_permute_idx, pre_all2all_inp_shape, input) + + post_all2all_fun = post_all2all(post_all2all_permute_idx, post_all2all_res_shape) + output = torch.empty_like(input_t) + dist.all_to_all_single(output, input_t, group=group) + + res = post_all2all_fun(output) + return res + + +class _SeqAllToAll(torch.autograd.Function): + + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + input: torch.Tensor, + scatter_idx: int, + gather_idx: int, + ) -> torch.Tensor: + ctx.group = group + ctx.scatter_idx = scatter_idx + ctx.gather_idx = gather_idx + res = single_all_to_all(input, scatter_idx, gather_idx, group) + return res + + @staticmethod + def backward(ctx: Any, *grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None, None]: + return None, _SeqAllToAll.apply(ctx.group, *grad_output, ctx.gather_idx, ctx.scatter_idx), None, None + + +class DistributedAttention(torch.nn.Module): + + def __init__( + self, + local_attention, + sequence_parallel, + scatter_idx: int = 2, + gather_idx: int = 1, + ) -> None: + super(DistributedAttention, self).__init__() + self.local_attn = local_attention + self.sequence_parallel = sequence_parallel + self.scatter_idx = scatter_idx + self.gather_idx = gather_idx + + def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attention_mask: torch.Tensor, *args: + Any, **kwargs) -> torch.Tensor: + if self.sequence_parallel.world_size == 1: + return self.local_attn(query, key, value, attention_mask, *args, **kwargs) + + # gather ulysses first, ring-attention next + if self.sequence_parallel.sp_world_size > 1: + query_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, query, self.scatter_idx, self.gather_idx) + key_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, key, self.scatter_idx, self.gather_idx) + value_layer = _SeqAllToAll.apply(self.sequence_parallel.sp_group, value, self.scatter_idx, self.gather_idx) + else: + query_layer, key_layer, value_layer = query, key, value + + if self.sequence_parallel.rp_world_size > 1: + # if need ring-attention + kwargs.pop('position_ids', None) + # Get the real position ids, this is filled by `sequence_parallel.prepare_inputs` + # real position ids is different from the position_ids when model uses mrope + position_ids = self.sequence_parallel.real_position_ids + # pad and split it by zigzag method + position_ids = self.sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids) + else: + # only ulysses + position_ids = kwargs.pop('position_ids') + if position_ids is not None: + # Reuse the generic gather path so 2D and 3D position_ids share the same SP behavior. + position_ids = self.sequence_parallel.gather(position_ids.contiguous(), dim=-1, position_ids=None) + + context_layer = self.local_attn( + query_layer, key_layer, value_layer, attention_mask, *args, position_ids=position_ids, **kwargs) + + if self.sequence_parallel.sp_world_size > 1: + output = _SeqAllToAll.apply(self.sequence_parallel.sp_group, context_layer, self.gather_idx, + self.scatter_idx) + else: + output = context_layer + + return output diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/utils.py b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e7569a8eca48c45c4c6de63d25a0e9bdb7178411 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/utils.py @@ -0,0 +1,171 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import math +import torch +import torch.distributed as dist +from torch.nn import CrossEntropyLoss +from torch.utils.data import Sampler +from typing import Any, Iterator + +from swift.dataloader import DataLoaderDispatcher +from .sequence_parallel import sequence_parallel + + +class GatherTensor(torch.autograd.Function): + """Gather tensor from sequence group (autograd supported)""" + + @staticmethod + def forward(ctx, tensor, dim=0, position_ids=None): + ctx.dim = dim + if position_ids is not None: + position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids) + ctx.position_ids = position_ids + return sequence_parallel.gather(tensor, dim=dim, position_ids=position_ids) + + @staticmethod + def backward(ctx, grad_output): + grad_input = sequence_parallel.split(grad_output, dim=ctx.dim, position_ids=ctx.position_ids) + return grad_input, None, None + + +class GatherLoss(torch.autograd.Function): + """Gather loss from sequence group""" + + @staticmethod + def forward(ctx, loss, labels, gather_idx=None, position_ids=None): + """ + Args: + loss: loss tensor after splitting + labels: labels tensor after splitting + gather_idx: gather the tensors on this dim + """ + # change from label.shape to loss, because label may be None + ctx.scatter_shape = loss.shape[gather_idx or 0] + ctx.gather_idx = gather_idx or 0 + if position_ids is not None: + position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids) + ctx.position_ids = position_ids + output = sequence_parallel.gather(loss, dim=ctx.gather_idx, position_ids=position_ids) + if labels is not None: + labels_output = sequence_parallel.gather(labels, dim=ctx.gather_idx, position_ids=position_ids) + else: + labels_output = None + return output, labels_output + + @staticmethod + def backward(ctx, *grad_output): + _grad = grad_output[0] * sequence_parallel.world_size + if sequence_parallel.rp_world_size > 1: + _grad = sequence_parallel.split(_grad, dim=ctx.gather_idx, position_ids=ctx.position_ids).contiguous() + else: + _grad = _grad.split( + ctx.scatter_shape, dim=ctx.gather_idx)[dist.get_rank(sequence_parallel.sp_group)].contiguous() + return _grad, None, None, None + + +class ChunkedCrossEntropyLoss(torch.autograd.Function): + + @staticmethod + def forward(ctx, logits, labels, chunk_size): + ctx.save_for_backward(logits, labels) + ctx.chunk_size = chunk_size + + losses = [] + for i in range(math.ceil(logits.shape[0] / chunk_size)): + l_start = i * chunk_size + l_end = min((i + 1) * chunk_size, logits.shape[0]) + logits_chunk = logits[l_start:l_end] + labels_chunk = labels[l_start:l_end] + loss_fct = CrossEntropyLoss(reduction='none') + loss_chunk = loss_fct(logits_chunk, labels_chunk) + losses.append(loss_chunk) + del logits_chunk + del labels_chunk + all_losses = torch.cat(losses) + return all_losses + + @staticmethod + def backward(ctx: Any, *grad_outputs: Any): + logits, labels = ctx.saved_tensors + chunk_size = ctx.chunk_size + + for i in range(math.ceil(logits.shape[0] / chunk_size)): + l_start = i * chunk_size + l_end = min((i + 1) * chunk_size, logits.shape[0]) + logits_chunk = logits[l_start:l_end].detach().requires_grad_(True) + labels_chunk = labels[l_start:l_end] + loss_fct = CrossEntropyLoss(reduction='none') + with torch.enable_grad(): + loss_chunk = loss_fct(logits_chunk, labels_chunk) + grad_output_chunk = grad_outputs[0][l_start:l_end] + _loss_chunk = (loss_chunk * grad_output_chunk).sum() + grad_chunk = torch.autograd.grad(_loss_chunk, logits_chunk, retain_graph=False)[0] + logits[l_start:l_end] = grad_chunk + + return logits, None, None + + +class SequenceParallelSampler(Sampler): + """Sampler for sequence parallel training""" + + def __init__(self, sp_instance, dataset, shuffle: bool = True, seed=None, round_up: bool = True) -> None: + self.sp_instance = sp_instance + rank = dist.get_rank(sp_instance.device_mesh['data'].get_group()) + world_size = sp_instance.device_mesh['data'].size() + self.rank = rank + self.world_size = world_size + + self.dataset = dataset + self.shuffle = shuffle + assert seed is not None + self.seed = seed + self.epoch = 0 + self.round_up = round_up + + if self.round_up: + self.num_samples = math.ceil(len(self.dataset) / world_size) + self.total_size = self.num_samples * self.world_size + else: + self.num_samples = math.ceil((len(self.dataset) - rank) / world_size) + self.total_size = len(self.dataset) + + def __iter__(self) -> Iterator[int]: + if self.shuffle: + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + indices = torch.randperm(len(self.dataset), generator=g).tolist() + else: + indices = torch.arange(len(self.dataset)).tolist() + + if self.round_up: + indices = (indices * int(self.total_size / len(indices) + 1))[:self.total_size] + + indices = indices[self.rank:self.total_size:self.world_size] + return iter(indices) + + def __len__(self) -> int: + return self.num_samples + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + +class SequenceParallelDispatcher(DataLoaderDispatcher): + """Dispatcher for sequence parallel training""" + + def __init__(self, dataloader, sp_instance, device=None, skip_batches: int = 0): + super().__init__(dataloader) + self.sp_instance = sp_instance + self.device = device + self.skip_batches = skip_batches + + @property + def rank(self): + return self.sp_instance.dp_rank if dist.is_initialized() else 0 + + @property + def world_size(self): + return self.sp_instance.dp_world_size if dist.is_initialized() else 1 + + @property + def group(self): + return self.sp_instance.dp_group if dist.is_initialized() else 1 diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/zigzag_ring_attn.py b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/zigzag_ring_attn.py new file mode 100644 index 0000000000000000000000000000000000000000..b9a60c716b8440f338085cd405de194855397527 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/zigzag_ring_attn.py @@ -0,0 +1,711 @@ +# Some code borrowed from the awesome work: https://github.com/zhuzilin/ring-flash-attention +# Copyright (c) ModelScope Contributors. All rights reserved. +import inspect +import torch +import torch.distributed as dist +import torch.nn.functional as F +from functools import cache + +from .ring_utils import RingComm +from .zigzag_ring_attn_npu import is_npu_tensor, npu_backward, npu_forward + + +def get_half_index(cu_seqlens, *, front: bool): + """Get half of the index + + Args: + cu_seqlens: The cu_seqlens passed into flash_attn + front: The head part or the tail part + + Returns: + The slice or the tensor mask. + """ + + if len(cu_seqlens) == 2: + if front: + return slice(None, cu_seqlens[-1] // 2) + else: + return slice(cu_seqlens[-1] // 2, None) + + index = torch.zeros((cu_seqlens[-1].item(), ), dtype=torch.bool) + for i in range(len(cu_seqlens) - 1): + start, end = cu_seqlens[i], cu_seqlens[i + 1] + if front: + end = (start + end) // 2 + else: + start = (start + end) // 2 + index[start:end] = True + return index + + +@torch.jit.script +def get_half_lse(lse, cu_seqlens, *, front: bool): + """Get half of the lse + + Args: + lse: The input lse, with shape [num_heads, seqlen] + cu_seqlens: The cu_seqlens passed into flash_attn + front: The head part or the tail part + + Returns: + The filtered lse with the same shape as lse + """ + new_lse = torch.empty( + (lse.shape[0], lse.shape[1] // 2), + dtype=lse.dtype, + device=lse.device, + ) + for i in range(len(cu_seqlens) - 1): + start, end = cu_seqlens[i].item(), cu_seqlens[i + 1].item() + new_start, new_end = start // 2, end // 2 + if front: + end -= (end - start) // 2 + else: + start += (end - start) // 2 + new_lse[:, new_start:new_end] = lse[:, start:end] + return new_lse + + +def update_out_and_lse(out, lse, block_out, block_lse): + """Update output and lse: + new_lse = lse + torch.log(1 + torch.exp(block_lse - lse)) + torch.exp(lse - new_lse) * out + torch.exp(block_lse - new_lse) * block_out + # For additional context and discussion, please refer to: + # https://github.com/zhuzilin/ring-flash-attention/pull/34#issuecomment-2076126795 + + Args: + out: The accumulated output of shape [seqlen, num_heads, hidden_size] + lse: The accumulated lse of shape [num_heads, seqlen] + block_out: The current block output of shape [seqlen, num_heads, hidden_size] + block_lse: The current block lse of shape [num_heads, seqlen] + + Returns: + The updated output[seqlen, num_heads, hidden_size] and lse (shape: [seqlen, num_heads, 1]) and + the intermediate value of torch.sigmoid(block_lse - lse) (shape: [seqlen, num_heads, 1]) + """ + if out is None: + out = block_out.to(torch.float32) + lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1) + sig_diff = None + else: + block_out = block_out.to(torch.float32) + block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1) + + diff = block_lse - lse + sig_diff = torch.sigmoid(diff) + + out = out - sig_diff * (out - block_out) # (..., D) + lse = lse - F.logsigmoid(lse - block_lse) # (..., 1) + return out, lse, sig_diff + + +@cache +def _get_default_args(func): + spec = inspect.getfullargspec(func) + defaults = spec.defaults if spec.defaults is not None else () + padded_defaults = (None, ) * (len(spec.args) - len(defaults)) + defaults + args = dict(zip(spec.args, padded_defaults)) + if 'softcap' in args: + args['softcap'] = 0.0 + return args + + +def get_default_args(func): + if inspect.isfunction(func): + return _get_default_args(func) + else: + # Use the origin _init_fn in CustomOpDef + return _get_default_args(func._init_fn) + + +def squeeze_batch(*t): + """Squeeze the batch dim of tensors""" + tensors = [] + for sub in t: + if sub.shape[0] == 1: + tensors.append(sub.squeeze(0)) + else: + tensors.append(sub) + return tuple(tensors) + + +def padding(tensor, cu_seqlens, padding_value, front): + """Pad the tensor according to the cu_seqlens + + Args: + tensor: The input tensor of shape [seqlen, *] + cu_seqlens: The cu_seqlens + padding_value: The padding value + front: tensor is the head or tail part + """ + if len(cu_seqlens) == 2: + if front: + return torch.cat((tensor, torch.full_like(tensor, padding_value).to(tensor.dtype).to(tensor.device)), dim=0) + else: + return torch.cat((torch.full_like(tensor, padding_value).to(tensor.dtype).to(tensor.device), tensor), dim=0) + + output = [] + acc = 0 + for i in range(len(cu_seqlens) - 1): + start, end = cu_seqlens[i], cu_seqlens[i + 1] + half_len = (end - start) // 2 + acc += half_len + half_start = start // 2 + local_tensor = tensor[half_start:half_start + half_len] + if front: + output.append(local_tensor) + output.append(torch.full_like(local_tensor, padding_value).to(local_tensor.dtype).to(local_tensor.device)) + else: + output.append(torch.full_like(local_tensor, padding_value).to(local_tensor.dtype).to(local_tensor.device)) + output.append(local_tensor) + assert acc == tensor.shape[0] + return torch.cat(output) + + +def forward(q, k, v, causal, cu_seqlens, max_seqlen, block_seq_len, dropout_p, softmax_scale, alibi_slopes, + window_size): + seqlen_q = q.shape[0] + seqlen_kv = k.shape[0] + half_cu_seqlens = cu_seqlens // 2 + half_max_seqlen = max_seqlen // 2 + cu_seqlens_q = half_cu_seqlens if seqlen_q == block_seq_len else cu_seqlens + max_seqlen_q = half_max_seqlen if seqlen_q == block_seq_len else max_seqlen + cu_seqlens_kv = half_cu_seqlens if seqlen_kv == block_seq_len else cu_seqlens + max_seqlen_kv = half_max_seqlen if seqlen_kv == block_seq_len else max_seqlen + if is_npu_tensor(q): + # Keep the ring schedule in this file unchanged; only the per-block + # flash-attn call is swapped to Ascend's TND varlen attention kernel. + return npu_forward( + q, + k, + v, + causal, + cu_seqlens_q, + cu_seqlens_kv, + dropout_p, + softmax_scale, + deterministic=False, + window_size=window_size, + ) + from flash_attn.flash_attn_interface import _flash_attn_varlen_forward + params = get_default_args(_flash_attn_varlen_forward).copy() + params.update({ + 'q': q, + 'k': k, + 'v': v, + # the first half and the second half are the same + 'cu_seqlens_q': cu_seqlens_q, + 'cu_seqlens_k': cu_seqlens_kv, + 'max_seqlen_q': max_seqlen_q, + 'max_seqlen_k': max_seqlen_kv, + 'dropout_p': dropout_p, + 'softmax_scale': softmax_scale, + 'causal': causal, + 'alibi_slopes': alibi_slopes, + 'return_softmax': True and dropout_p > 0, + }) + if 'window_size' in params: + params.update({'window_size': window_size}) + else: + params.update({ + 'window_size_left': window_size[0], + 'window_size_right': window_size[1], + }) + assert k.shape[-0] == cu_seqlens_kv[-1] + assert q.shape[-0] == cu_seqlens_q[-1] + assert max_seqlen_q == (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item() + assert max_seqlen_kv == (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]).max().item() + outputs = _flash_attn_varlen_forward(**params) + if len(outputs) == 8: + block_out, _, _, _, _, block_lse, _, _ = outputs + else: + assert len(outputs) == 4 + block_out, block_lse, _, _ = outputs + return block_out, block_lse + + +def backward(dout, q, k, v, out, softmax_lse, causal, cu_seqlens, max_seqlen, block_seq_len, dq_buffer, dk_buffer, + dv_buffer, dropout_p, softmax_scale, alibi_slopes, deterministic, window_size): + seqlen_q = q.shape[0] + seqlen_kv = k.shape[0] + + half_cu_seqlens = cu_seqlens // 2 + half_max_seqlen = max_seqlen // 2 + cu_seqlens_q = half_cu_seqlens if seqlen_q == block_seq_len else cu_seqlens + max_seqlen_q = half_max_seqlen if seqlen_q == block_seq_len else max_seqlen + cu_seqlens_kv = half_cu_seqlens if seqlen_kv == block_seq_len else cu_seqlens + max_seqlen_kv = half_max_seqlen if seqlen_kv == block_seq_len else max_seqlen + from flash_attn.flash_attn_interface import _flash_attn_varlen_backward + params = get_default_args(_flash_attn_varlen_backward).copy() + params.update({ + 'dout': dout, + 'q': q, + 'k': k, + 'v': v, + 'out': out, + 'softmax_lse': softmax_lse, + 'dq': dq_buffer[:seqlen_q], + 'dk': dk_buffer[:seqlen_kv], + 'dv': dv_buffer[:seqlen_kv], + # the first half and the second half are the same + 'cu_seqlens_q': cu_seqlens_q, + 'cu_seqlens_k': cu_seqlens_kv, + 'max_seqlen_q': max_seqlen_q, + 'max_seqlen_k': max_seqlen_kv, + 'dropout_p': dropout_p, + 'softmax_scale': softmax_scale, + 'causal': causal, + 'alibi_slopes': alibi_slopes, + 'deterministic': deterministic, + }) + assert dout.shape[0] == q.shape[0] + assert dout.shape[0] == out.shape[0] + assert softmax_lse.shape[1] == q.shape[0] + assert k.shape[0] == cu_seqlens_kv[-1] + assert q.shape[0] == cu_seqlens_q[-1] + assert max_seqlen_q == (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).max().item() + assert max_seqlen_kv == (cu_seqlens_kv[1:] - cu_seqlens_kv[:-1]).max().item() + if 'window_size' in params: + params.update({'window_size': window_size}) + else: + params.update({ + 'window_size_left': window_size[0], + 'window_size_right': window_size[1], + }) + _flash_attn_varlen_backward(**params) + + +def lse_grad(out, lse, block_out, block_lse, sig, grad_out, grad_lse): + """Calculate the grad of each block. + + Args: + out: The accumulated output of shape [seqlen, num_heads, hidden_size] + lse: The accumulated lse of shape [num_heads, seqlen, 1] + block_out: The current block output of shape [seqlen, num_heads, hidden_size] + block_lse: The current block lse of shape [num_heads, seqlen, 1] + grad_out: The input grad of output of the current block shape [seqlen, num_heads, hidden_size] + grad_lse: The input grad of lse of the current block shape [num_heads, seqlen, 1] + + Returns: + The accumulated grad of out and lse, and the grad of out and lse of the current block + """ + grad_out_input = grad_out * (1 - sig) + + grad_block_out = grad_out * sig + + d_new_out_d_lse = (out - block_out) * (sig * (1 - sig)) + grad_lse_input = (grad_out * d_new_out_d_lse).sum(dim=-1, keepdim=True) + grad_lse_input_final = grad_lse_input + grad_lse * torch.sigmoid(lse - block_lse) + + grad_block_lse = -grad_lse_input_final + grad_lse + + return grad_out_input, grad_lse_input_final, grad_block_out, grad_block_lse + + +def zigzag_ring_flash_attn_varlen_forward( + process_group, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens, + max_seqlen, + half_index0, + half_index1, + softmax_scale, + dropout_p=0, + causal=True, + window_size=(-1, -1), + alibi_slopes=None, + deterministic=False, +): + assert causal, 'zigzag ring is meaningless for causal=False' + comm = RingComm(process_group) + q, k, v = squeeze_batch(q, k, v) + q1 = q[half_index1] + # Input cu_seqlens is the total length, divided by world_size to fit the split ones + cu_seqlens = cu_seqlens // comm.world_size + # Same with above + max_seqlen = max_seqlen // comm.world_size + block_seq_len = q.shape[0] // 2 + out = None + lse = None + next_k, next_v = None, None + + for step in range(comm.world_size): + # from step 0 to the last + if step + 1 != comm.world_size: + next_k, next_v = comm.send_recv_kv(k, v) + """ + world_size = 4, total 8 parts + 0/7 is group0 + 1/6 is group1 + 2/5 is group2 + 3/4 is group3 + consider 1/6，take the query as the left axis, key as the top axis: + step 0: + 1 6 + 1 ✅ ❎ + 6 ✅ ✅ + all needed, causal=True + step 1(step <= comm.rank): + 0 7 + 1 ✅ ❎ + 6 ✅ ❎ + the first part of kv is needed, causal=False + step 2(step > comm.rank): + 3 4 + 1 ❎ ❎ + 6 ✅ ✅ + the second part of q is needed, causal=False + """ + # Here block_lse shape: [num_heads, seqlen] + # lse shape: [seqlen, num_heads, 1] + if step == 0: + block_out, block_lse = forward(q, k, v, True, cu_seqlens, max_seqlen, block_seq_len, dropout_p, + softmax_scale, alibi_slopes, window_size) + out, lse, sig_diff = update_out_and_lse(out, lse, block_out, block_lse) + elif step <= comm.rank: + k0 = k[half_index0] + v0 = v[half_index0] + block_out, block_lse = forward(q, k0, v0, False, cu_seqlens, max_seqlen, block_seq_len, dropout_p, + softmax_scale, alibi_slopes, window_size) + out, lse, sig_diff = update_out_and_lse(out, lse, block_out, block_lse) + else: + block_out, block_lse = forward(q1, k, v, False, cu_seqlens, max_seqlen, block_seq_len, dropout_p, + softmax_scale, alibi_slopes, window_size) + out[half_index1], lse[half_index1], sig_diff = update_out_and_lse(out[half_index1], lse[half_index1], + block_out, block_lse) + + if step + 1 != comm.world_size: + comm.wait() + k, v = next_k, next_v + + out = out.to(q.dtype) + lse = lse.squeeze(dim=-1).transpose(0, 1) # [num_heads, seqlen] + return out.unsqueeze(0), lse.unsqueeze(0) + + +def zigzag_ring_flash_attn_varlen_backward( + process_group, + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens, + max_seqlen, + half_index0, + half_index1, + softmax_scale, + dropout_p=0, + causal=True, + window_size=(-1, -1), + alibi_slopes=None, + deterministic=False, +): + assert causal, 'zigzag ring is meaningless for causal=False' + if is_npu_tensor(q): + # NPU backward uses native flash-attn grad with the final ring out/lse + # patched into each block ctx. Missing kernel support should fail loudly. + return npu_backward( + process_group, + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens, + max_seqlen, + half_index0, + half_index1, + softmax_scale, + dropout_p=dropout_p, + window_size=window_size, + deterministic=deterministic, + ) + kv_comm = RingComm(process_group) + d_kv_comm = RingComm(process_group) + dk_comm_buffer = dv_comm_buffer = None + dq, dk, dv = None, None, None + next_dk, next_dv = None, None + next_k, next_v = None, None + + # squeeze the axis of batch + dout, q, k, v, out, softmax_lse = squeeze_batch(dout, q, k, v, out, softmax_lse) + q1 = q[half_index1] + # Input cu_seqlens is the total length, divided by world_size to fit the split ones + cu_seqlens = cu_seqlens // kv_comm.world_size + # Same as above + max_seqlen = max_seqlen // kv_comm.world_size + # half of the part + block_seq_len = q.shape[0] // 2 + + # repeatly allocating buffer may be slow... + dq_buffer = torch.empty(q.shape, dtype=q.dtype, device=q.device) + dk_buffer = torch.empty(k.shape, dtype=k.dtype, device=k.device) + dv_buffer = torch.empty(v.shape, dtype=v.dtype, device=v.device) + origin_q, origin_k, origin_v = q, k, v + + out_lse = [] + fout = None + flse = None + # Recalculate forward with the same qkv to generate out_lse, used to calculate the grad + for step in range(kv_comm.world_size): + if step + 1 != kv_comm.world_size: + next_k, next_v = kv_comm.send_recv_kv(k, v) + + if step == 0: + block_out, block_lse = forward(q, k, v, True, cu_seqlens, max_seqlen, block_seq_len, dropout_p, + softmax_scale, alibi_slopes, window_size) + fout, flse, sig_diff = update_out_and_lse(fout, flse, block_out, block_lse) + elif step <= kv_comm.rank: + k0 = k[half_index0] + v0 = v[half_index0] + block_out, block_lse = forward(q, k0, v0, False, cu_seqlens, max_seqlen, block_seq_len, dropout_p, + softmax_scale, alibi_slopes, window_size) + fout, flse, sig_diff = update_out_and_lse(fout, flse, block_out, block_lse) + else: + block_out, block_lse = forward(q1, k, v, False, cu_seqlens, max_seqlen, block_seq_len, dropout_p, + softmax_scale, alibi_slopes, window_size) + fout[half_index1], flse[half_index1], sig_diff = update_out_and_lse(fout[half_index1], flse[half_index1], + block_out, block_lse) + + block_lse = block_lse.transpose(0, 1).unsqueeze(-1) + if step > kv_comm.rank: + # cat zeros because there are may be a half of the out/lse + block_out = padding(block_out, cu_seqlens, 0, front=False) + block_lse = padding(block_lse, cu_seqlens, -1e5, front=False) + sig_diff = padding(sig_diff, cu_seqlens, 0, front=False) + + # save to out_lse + out_lse.append((fout, flse, block_out, block_lse, sig_diff)) + + if step + 1 != kv_comm.world_size: + kv_comm.wait() + k, v = next_k, next_v + + current_dout = dout + current_dlse = torch.zeros_like(softmax_lse.transpose(0, 1).unsqueeze(-1)) + block_gradients = {} + + for i in reversed(range(len(out_lse))): + if i == 0: + # the first step does not need + continue + stored_out, stored_lse, stored_block_out, stored_block_lse, stored_sig = out_lse[i] + grad_out_input, grad_lse_input, grad_block_out, grad_block_lse = lse_grad(stored_out, stored_lse, + stored_block_out, stored_block_lse, + stored_sig, current_dout, + current_dlse) + current_dout = grad_out_input + current_dlse = grad_lse_input + block_gradients[i] = {'grad_block_out': grad_block_out, 'grad_block_lse': grad_block_lse} + + q, k, v = origin_q, origin_k, origin_v + + for step in range(kv_comm.world_size): + _, _, block_out, block_lse, _ = out_lse[step] + if block_out.isnan().any() or block_lse.isnan().any(): + raise + block_lse = block_lse.transpose(0, 1).squeeze(2) + + if step + 1 != kv_comm.world_size: + next_k, next_v = kv_comm.send_recv_kv(k, v) + + if step == 0: + # if step == 0, use the final current_dout + block_dout = current_dout + else: + # else use the grad in the block_gradients + block_dout = block_gradients[step]['grad_block_out'] + + if block_dout.isnan().any(): + raise + + if step == 0: + backward( + block_dout.to(dout.dtype), q, k, v, block_out, block_lse, True, cu_seqlens, max_seqlen, block_seq_len, + dq_buffer, dk_buffer, dv_buffer, dropout_p, softmax_scale, alibi_slopes, deterministic, window_size) + dq = dq_buffer.to(torch.float32) + dk = dk_buffer.to(torch.float32) + dv = dv_buffer.to(torch.float32) + if dq.isnan().any() or dk.isnan().any() or dv.isnan().any(): + raise + else: + if step <= kv_comm.rank: + k0 = k[half_index0] + v0 = v[half_index0] + backward( + block_dout.to(dout.dtype), q, k0, v0, block_out, block_lse, False, cu_seqlens, max_seqlen, + block_seq_len, dq_buffer, dk_buffer, dv_buffer, dropout_p, softmax_scale, alibi_slopes, + deterministic, window_size) + dq += dq_buffer + else: + backward(block_dout[half_index1].to(dout.dtype), q1, k, v, block_out[half_index1], + get_half_lse(block_lse, cu_seqlens, + front=False), False, cu_seqlens, max_seqlen, block_seq_len, dq_buffer, dk_buffer, + dv_buffer, dropout_p, softmax_scale, alibi_slopes, deterministic, window_size) + # only need to add to the tail half, because the head half does not match the causal condition + dq[half_index1] += dq_buffer[:block_seq_len] + + d_kv_comm.wait() + # dk_comm_buffer, dv_comm_buffer = dk, dv + # avoid d_kv_comm.send_recv_kv causing dk_comm_buffer reuse the same memory with next_dk and dk + dk_comm_buffer = torch.empty_like(dk) + dv_comm_buffer = torch.empty_like(dv) + dk_comm_buffer.copy_(dk) + dv_comm_buffer.copy_(dv) + # next_dk, next_dv comes from a previous gpu, add kv grad to them, and pass them to the next gpu + dk, dv = next_dk, next_dv + + if step <= kv_comm.rank: + # only need to add to the head part, because the tail part does not match the causal condition + dk[half_index0] += dk_buffer[:block_seq_len] + dv[half_index0] += dv_buffer[:block_seq_len] + else: + dk += dk_buffer + dv += dv_buffer + if dq.isnan().any() or dk.isnan().any() or dv.isnan().any(): + raise + if step + 1 != kv_comm.world_size: + kv_comm.wait() + k, v = next_k, next_v + + next_dk, next_dv = d_kv_comm.send_recv_kv(dk, dv, dk_comm_buffer, dv_comm_buffer) + + d_kv_comm.wait() + + return dq.to(q.dtype).unsqueeze(0), next_dk.to(q.dtype).unsqueeze(0), next_dv.to(q.dtype).unsqueeze(0) + + +class ZigZagRingFlashAttnVarlenFunc(torch.autograd.Function): + + @staticmethod + def forward( + ctx, + q, + k, + v, + cu_seqlens, + max_seqlen, + dropout_p, + softmax_scale, + causal, + window_size, + alibi_slopes, + deterministic, + return_softmax, + group, + ): + if softmax_scale is None: + softmax_scale = q.shape[-1]**(-0.5) + + assert alibi_slopes is None + k = k.contiguous() + v = v.contiguous() + rp_world_size = dist.get_world_size(group) + half_index0 = get_half_index(cu_seqlens // rp_world_size, front=True) + half_index1 = get_half_index(cu_seqlens // rp_world_size, front=False) + out, softmax_lse = zigzag_ring_flash_attn_varlen_forward( + group, + q, + k, + v, + cu_seqlens, + max_seqlen, + half_index0, + half_index1, + softmax_scale=softmax_scale, + dropout_p=dropout_p, + causal=causal, + window_size=window_size, + alibi_slopes=alibi_slopes, + deterministic=False, + ) + # this should be out_padded + is_half_index_tensor = isinstance(half_index0, torch.Tensor) + ctx.is_half_index_tensor = is_half_index_tensor + if is_half_index_tensor: + """ + Shapes: + qkv: [1, seqlen, num_heads, hidden_size] + out: [1, seqlen, num_heads, hidden_size] + softmax_lse: [1, num_heads, seqlen] + """ + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens, half_index0, half_index1) + else: + ctx.save_for_backward(q, k, v, out, softmax_lse, cu_seqlens) + ctx.half_index0 = half_index0 + ctx.half_index1 = half_index1 + ctx.max_seqlen = max_seqlen + ctx.dropout_p = dropout_p + ctx.softmax_scale = softmax_scale + ctx.causal = causal + ctx.window_size = window_size + ctx.alibi_slopes = alibi_slopes + ctx.deterministic = deterministic + ctx.group = group + return out if not return_softmax else (out, softmax_lse, None) + + @staticmethod + def backward(ctx, dout, *args): + if ctx.is_half_index_tensor: + (q, k, v, out, softmax_lse, cu_seqlens, half_index0, half_index1) = (ctx.saved_tensors) + else: + q, k, v, out, softmax_lse, cu_seqlens = ctx.saved_tensors + half_index0 = ctx.half_index0 + half_index1 = ctx.half_index1 + dq, dk, dv = zigzag_ring_flash_attn_varlen_backward( + ctx.group, + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens, + ctx.max_seqlen, + half_index0, + half_index1, + softmax_scale=ctx.softmax_scale, + dropout_p=ctx.dropout_p, + causal=ctx.causal, + window_size=ctx.window_size, + alibi_slopes=ctx.alibi_slopes, + deterministic=ctx.deterministic, + ) + return dq, dk, dv, None, None, None, None, None, None, None, None, None, None + + +def zigzag_ring_flash_attn_varlen_func( + q, + k, + v, + cu_seqlens, + max_seqlen, + dropout_p=0.0, + softmax_scale=None, + causal=False, + window_size=(-1, -1), # -1 means infinite context window + alibi_slopes=None, + deterministic=False, + return_attn_probs=False, + group=None, +): + return ZigZagRingFlashAttnVarlenFunc.apply( + q, + k, + v, + cu_seqlens, + max_seqlen, + dropout_p, + softmax_scale, + causal, + window_size, + alibi_slopes, + deterministic, + return_attn_probs, + group, + ) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/zigzag_ring_attn_npu.py b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/zigzag_ring_attn_npu.py new file mode 100644 index 0000000000000000000000000000000000000000..965dc42492f2a3e1c1e77a92175e1aa03bf82635 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/sequence_parallel/zigzag_ring_attn_npu.py @@ -0,0 +1,459 @@ +# Some code borrowed from the awesome work: https://github.com/zhuzilin/ring-flash-attention +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from functools import cache + +from .ring_utils import RingComm + +_NPU_BLOCK_MASK_SIZE = 2048 +_NPU_FULL_TOKENS = 2147483647 +_NPU_TND_SOFTMAX_STAT_REPEAT = 8 + + +def is_npu_tensor(tensor: torch.Tensor) -> bool: + return tensor.device.type == 'npu' + + +def _cu_seqlens_to_actual_seq(cu_seqlens: torch.Tensor) -> tuple[int, ...]: + return tuple(int(x) for x in cu_seqlens[1:].detach().cpu().tolist()) + + +@cache +def _get_npu_causal_mask_cpu() -> torch.Tensor: + return torch.triu(torch.ones((_NPU_BLOCK_MASK_SIZE, _NPU_BLOCK_MASK_SIZE), dtype=torch.bool), diagonal=1) + + +def _get_npu_causal_mask(device: torch.device) -> torch.Tensor: + return _get_npu_causal_mask_cpu().to(device=device) + + +def _normalize_window_size(window_size): + if window_size is None: + return -1, -1 + return window_size + + +def _get_npu_sparse_params(causal: bool, window_size, device: torch.device) -> dict: + window_size = _normalize_window_size(window_size) + if window_size != (-1, -1): + left, right = window_size + left = _NPU_FULL_TOKENS if left < 0 else int(left) + right = _NPU_FULL_TOKENS if right < 0 else int(right) + if causal: + right = 0 + return { + 'atten_mask': _get_npu_causal_mask(device), + 'sparse_mode': 4, + 'pre_tockens': left, + 'next_tockens': right, + } + if causal: + return { + 'atten_mask': _get_npu_causal_mask(device), + 'sparse_mode': 3, + 'pre_tockens': _NPU_FULL_TOKENS, + 'next_tockens': _NPU_FULL_TOKENS, + } + return { + 'atten_mask': None, + 'sparse_mode': 0, + 'pre_tockens': _NPU_FULL_TOKENS, + 'next_tockens': _NPU_FULL_TOKENS, + } + + +def _reshape_npu_lse(lse: torch.Tensor, seqlen_q: int, num_heads: int) -> torch.Tensor: + """Normalize Ascend softmax stats to flash-attn's [num_heads, seqlen] layout.""" + if lse.dim() == 2: + if lse.shape == (num_heads, seqlen_q): + return lse.contiguous() + if lse.shape == (seqlen_q, num_heads): + return lse.transpose(0, 1).contiguous() + elif lse.dim() == 3: + # Some CANN versions return an extra trailing size-8 axis with repeated + # stats. Ring merge only needs one copy of each token/head lse. + if lse.shape[-1] == 8: + lse = lse[..., 0] + if lse.shape == (seqlen_q, num_heads): + return lse.transpose(0, 1).contiguous() + if lse.shape == (num_heads, seqlen_q): + return lse.contiguous() + if lse.shape[0] == seqlen_q: + return lse.permute(1, 2, 0).reshape(num_heads, seqlen_q).contiguous() + if lse.shape[1] == seqlen_q: + return lse.permute(0, 2, 1).reshape(num_heads, seqlen_q).contiguous() + raise RuntimeError(f'Unexpected NPU lse shape {tuple(lse.shape)} for seqlen_q={seqlen_q}, num_heads={num_heads}') + + +def _get_npu_attention_common_kwargs( + q: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + softmax_scale: float, + dropout_p: float, + causal: bool, + window_size, + deterministic: bool, +) -> dict: + sparse_params = _get_npu_sparse_params(causal, window_size, q.device) + return { + 'head_num': q.shape[1], + 'input_layout': 'TND', + 'scale_value': softmax_scale or q.shape[-1]**(-0.5), + 'keep_prob': 1. - dropout_p, + 'actual_seq_qlen': _cu_seqlens_to_actual_seq(cu_seqlens_q), + 'actual_seq_kvlen': _cu_seqlens_to_actual_seq(cu_seqlens_kv), + 'sync': bool(deterministic and dropout_p > 0), + **sparse_params, + } + + +def _call_npu_fusion_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + softmax_scale: float, + dropout_p: float, + causal: bool, + window_size, + deterministic: bool, +): + import torch_npu + + common_kwargs = _get_npu_attention_common_kwargs( + q, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + softmax_scale=softmax_scale, + dropout_p=dropout_p, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ) + params = { + 'query': q, + 'key': k, + 'value': v, + 'scale': common_kwargs['scale_value'], + 'softmax_layout': 'TND', + } + params.update(common_kwargs) + params.pop('scale_value') + return torch_npu.npu_fusion_attention(**params) + + +def _call_npu_fusion_attention_grad( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + attention_out: torch.Tensor, + softmax_lse: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + softmax_scale: float, + dropout_p: float, + causal: bool, + window_size, + deterministic: bool, +): + import torch_npu + + if not hasattr(torch_npu, 'npu_fusion_attention_grad'): + raise AttributeError('torch_npu.npu_fusion_attention_grad is not available') + # Dropout backward needs the exact seed/offset from the original forward, + # which this ring ctx does not save. Fail instead of using a wrong mask. + if dropout_p != 0.0: + raise NotImplementedError('NPU ring attention native backward currently requires dropout_p=0.') + + common_kwargs = _get_npu_attention_common_kwargs( + q, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + softmax_scale=softmax_scale, + dropout_p=dropout_p, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ) + softmax_max, softmax_sum = _npu_softmax_stats_from_global_lse( + softmax_lse, + q_tokens=q.shape[0], + num_heads=q.shape[1], + ) + + params = { + 'query': q, + 'key': k, + 'value': v, + 'dy': dout, + 'head_num': common_kwargs['head_num'], + 'input_layout': common_kwargs['input_layout'], + 'atten_mask': common_kwargs['atten_mask'], + 'softmax_max': softmax_max, + 'softmax_sum': softmax_sum, + 'softmax_in': None, + 'attention_in': (attention_out if torch.is_tensor(attention_out) and attention_out.numel() > 0 else None), + 'scale_value': common_kwargs['scale_value'], + 'keep_prob': common_kwargs['keep_prob'], + 'pre_tockens': common_kwargs['pre_tockens'], + 'next_tockens': common_kwargs['next_tockens'], + 'seed': 0, + 'offset': 0, + 'numels': 0, + 'actual_seq_qlen': common_kwargs['actual_seq_qlen'], + 'actual_seq_kvlen': common_kwargs['actual_seq_kvlen'], + 'sparse_mode': common_kwargs['sparse_mode'], + 'sync': common_kwargs['sync'], + 'softmax_layout': 'TND', + } + return torch_npu.npu_fusion_attention_grad(**params) + + +def _normalize_flash_attn_lse(softmax_lse: torch.Tensor, total_len: int) -> torch.Tensor: + """Normalize flash-attn lse to [num_heads, total_len].""" + lse = softmax_lse + if lse.dim() == 3 and lse.shape[0] == 1: + lse = lse.squeeze(0) + if lse.dim() != 2: + raise RuntimeError(f'Unexpected softmax_lse shape: {tuple(softmax_lse.shape)}') + if lse.shape[1] != total_len: + lse = lse.transpose(0, 1).contiguous() + if lse.shape[1] != total_len: + raise RuntimeError(f'Unexpected softmax_lse shape: {tuple(softmax_lse.shape)} for total_len={total_len}') + return lse + + +def _npu_softmax_stats_from_global_lse( + softmax_lse_global: torch.Tensor, + q_tokens: int, + num_heads: int, +) -> tuple[torch.Tensor, torch.Tensor]: + lse_h_t = _normalize_flash_attn_lse(softmax_lse_global, q_tokens) + if lse_h_t.shape[0] != num_heads: + raise RuntimeError(f'Unexpected global lse shape: {tuple(softmax_lse_global.shape)} ' + f'for q_tokens={q_tokens}, num_heads={num_heads}') + + # With softmax_layout='TND', Ascend returns softmax stats as [T, N, 8]. + # The split-attention backward only needs logsumexp; max=lse and sum=1 + # encode the same value without replaying the block forward. + lse_t_h = lse_h_t.transpose(0, 1).contiguous().to(torch.float32) + softmax_max = lse_t_h.unsqueeze(-1).expand( + q_tokens, + num_heads, + _NPU_TND_SOFTMAX_STAT_REPEAT, + ).contiguous() + return softmax_max, torch.ones_like(softmax_max) + + +def _get_second_half_lse(softmax_lse: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: + total_len = int(cu_seqlens[-1].item()) + lse = _normalize_flash_attn_lse(softmax_lse, total_len) + + # The step > rank branch only differentiates q[half_index1]. Slice the final + # merged lse per sequence so the native grad ctx sees the same query span. + second_half_lse = torch.empty((lse.shape[0], lse.shape[1] // 2), dtype=lse.dtype, device=lse.device) + for i in range(len(cu_seqlens) - 1): + start, end = cu_seqlens[i].item(), cu_seqlens[i + 1].item() + new_start, new_end = start // 2, end // 2 + start += (end - start) // 2 + second_half_lse[:, new_start:new_end] = lse[:, start:end] + return second_half_lse + + +def _npu_block_backward_with_global_stats( + block_dout, + block_q, + block_k, + block_v, + block_out_global, + block_lse_global, + block_causal, + cu_seqlens_q, + cu_seqlens_kv, + softmax_scale, + dropout_p, + window_size, + deterministic, +): + """Run one native NPU block backward using the final merged ring stats.""" + return _call_npu_fusion_attention_grad( + block_dout.to(block_q.dtype), + block_q, + block_k, + block_v, + attention_out=block_out_global, + softmax_lse=block_lse_global, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + softmax_scale=softmax_scale, + dropout_p=dropout_p, + causal=block_causal, + window_size=window_size, + deterministic=deterministic, + )[:3] + + +def _squeeze_batch(*tensors): + squeezed = [] + for tensor in tensors: + if tensor.shape[0] == 1: + squeezed.append(tensor.squeeze(0)) + else: + squeezed.append(tensor) + return tuple(squeezed) + + +def npu_backward( + process_group, + dout, + q, + k, + v, + out, + softmax_lse, + cu_seqlens, + max_seqlen, + half_index0, + half_index1, + softmax_scale, + dropout_p=0.0, + window_size=(-1, -1), + deterministic=False, +): + kv_comm = RingComm(process_group) + d_kv_comm = RingComm(process_group) + dout, q, k, v, out, softmax_lse = _squeeze_batch(dout, q, k, v, out, softmax_lse) + cu_seqlens = cu_seqlens // kv_comm.world_size + del max_seqlen + + half_cu_seqlens = cu_seqlens // 2 + q1 = q[half_index1] + dout1 = dout[half_index1] + out1 = out[half_index1] + softmax_lse1 = _get_second_half_lse(softmax_lse, cu_seqlens) + + dq = torch.zeros_like(q, dtype=torch.float32) + current_step_dk = torch.empty_like(k, dtype=torch.float32) + current_step_dv = torch.empty_like(v, dtype=torch.float32) + next_dk = next_dv = None + + for step in range(kv_comm.world_size): + current_step_dk.zero_() + current_step_dv.zero_() + if step == 0: + bdq, bdk, bdv = _npu_block_backward_with_global_stats( + dout, + q, + k, + v, + out, + softmax_lse, + True, + cu_seqlens, + cu_seqlens, + softmax_scale, + dropout_p, + window_size, + deterministic, + ) + dq += bdq.to(torch.float32) + current_step_dk += bdk.to(torch.float32) + current_step_dv += bdv.to(torch.float32) + elif step <= kv_comm.rank: + k0 = k[half_index0] + v0 = v[half_index0] + bdq, bdk, bdv = _npu_block_backward_with_global_stats( + dout, + q, + k0, + v0, + out, + softmax_lse, + False, + cu_seqlens, + half_cu_seqlens, + softmax_scale, + dropout_p, + window_size, + deterministic, + ) + dq += bdq.to(torch.float32) + current_step_dk[half_index0] += bdk.to(torch.float32) + current_step_dv[half_index0] += bdv.to(torch.float32) + else: + bdq, bdk, bdv = _npu_block_backward_with_global_stats( + dout1, + q1, + k, + v, + out1, + softmax_lse1, + False, + half_cu_seqlens, + cu_seqlens, + softmax_scale, + dropout_p, + window_size, + deterministic, + ) + dq[half_index1] += bdq.to(torch.float32) + current_step_dk += bdk.to(torch.float32) + current_step_dv += bdv.to(torch.float32) + + # K/V gradients are owned by the rank that originally held that shard. + # Rotate the accumulated gradients in the opposite ring direction until + # each owner receives its final dk/dv. + if step == 0: + dk = current_step_dk + dv = current_step_dv + else: + dk = next_dk + dv = next_dv + dk += current_step_dk + dv += current_step_dv + + next_dk, next_dv = d_kv_comm.send_recv_kv(dk, dv) + d_kv_comm.wait() + + if step + 1 != kv_comm.world_size: + next_k, next_v = kv_comm.send_recv_kv(k, v) + kv_comm.wait() + k, v = next_k, next_v + + return dq.to(q.dtype).unsqueeze(0), next_dk.to(q.dtype).unsqueeze(0), next_dv.to(q.dtype).unsqueeze(0) + + +def npu_forward( + q, + k, + v, + causal, + cu_seqlens_q, + cu_seqlens_kv, + dropout_p, + softmax_scale, + deterministic, + window_size, +): + outputs = _call_npu_fusion_attention( + q, + k, + v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + softmax_scale=softmax_scale, + dropout_p=dropout_p, + causal=causal, + window_size=window_size, + deterministic=deterministic, + ) + block_out, softmax_max, softmax_sum = outputs[:3] + block_lse = softmax_max.to(torch.float32) + torch.log(softmax_sum.to(torch.float32)) + block_lse = _reshape_npu_lse(block_lse, q.shape[0], q.shape[1]) + return block_out, block_lse diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/__init__.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d1e0381ea9ceba75ad69eea2fadc5030be5be20 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +from . import templates +from .base import MaxLengthError, Template +from .constant import TemplateType +from .grounding import draw_bbox +from .register import TEMPLATE_MAPPING, get_template, get_template_meta, register_template +from .template_inputs import StdTemplateInputs, TemplateInputs +from .template_meta import TemplateMeta +from .utils import (ContextType, History, Messages, Prompt, Tool, Word, get_last_user_round, history_to_messages, + messages_to_history, split_str_parts_by, update_generation_config_eos_token) +from .vision_utils import load_file, load_image diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/base.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/base.py new file mode 100644 index 0000000000000000000000000000000000000000..3df77b97bec1a9d1d04e6b2ea9143a61f8730fd5 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/base.py @@ -0,0 +1,2347 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import hashlib +import inspect +import math +import numpy as np +import os +import random +import re +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections import defaultdict +from contextlib import contextmanager, nullcontext +from copy import deepcopy +from dataclasses import asdict +from functools import partial, wraps +from modelscope.hub.utils.utils import get_cache_dir +from peft import PeftModel +from PIL import Image +from torch.nn.utils.rnn import pad_sequence +from transformers import StoppingCriteriaList +from transformers.integrations import is_deepspeed_zero3_enabled +from transformers.utils import strtobool +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union + +from swift.utils import Processor, ProcessorMixin, get_env_args, get_logger, remove_response, retry_decorator, to_device +from .template_inputs import StdTemplateInputs, TemplateInputs +from .utils import Context, ContextType, StopWordsCriteria, fetch_one, findall, get_last_user_round, split_str_parts_by +from .vision_utils import _check_path, load_audio, load_batch, load_image, rescale_image + +logger = get_logger() +if TYPE_CHECKING: + from swift.infer_engine import InferRequest + from .template_meta import TemplateMeta + + +class MaxLengthError(ValueError): + pass + + +class Template(ProcessorMixin): + """Base template class for formatting and processing model inputs/outputs. + + This class serves as the foundation for all template implementations in the Swift framework. + It handles the conversion between conversation formats and token sequences, manages multimodal + inputs (images, videos, audio), supports various training modes (standard, RLHF, KTO), and + provides utilities for tokenization, padding, and data collation. + + The Template class is designed to be flexible and extensible, supporting: + - Multiple chat formats (user/assistant conversations, system prompts, tool calls) + - Multimodal data processing (images, videos, audio, bounding boxes) + - Different training strategies (causal language modeling, sequence classification, embedding, etc.) + - Various inference engines (Transformers, vLLM, LMDeploy, SGLang) + - Advanced features like padding-free training, sequence parallelism, and loss scaling + """ + special_tokens = ['', ' in the tool instruction + system = system.replace( + f'you should print ', + f'you should print <｜reasoning_mode{self.HYTK}｜>' + f'reasoning_effort:{reasoning_effort}') + else: + # For non-tool calls, append reasoning_mode to the system/prefix area + mode_str = f'<｜reasoning_mode{self.HYTK}｜>reasoning_effort:{reasoning_effort}' + system = (system or '') + mode_str + return system + + +register_template( + TemplateMeta( + LLMTemplateType.hy_v3_preview, + prefix=['<｜hy_begin▁of▁sentence｜>'], + system_prefix=['<｜hy_begin▁of▁sentence｜>{{SYSTEM}}'], + prompt=['<｜hy_User｜>{{QUERY}}<｜hy_Assistant｜>'], + chat_sep=['<｜hy_eos｜>'], + suffix=['<｜hy_eos｜>'], + template_cls=HyV3PreviewTemplate, + is_thinking=True, + thinking_prefix='', + non_thinking_prefix='', + history_thinking_prefix='', + agent_template='hy_v3_preview')) + + +class HyV3Template(HyV3PreviewTemplate): + HYTK = ':opensource' + + +register_template( + TemplateMeta( + LLMTemplateType.hy_v3, + prefix=['<｜hy_begin_of_sentence:opensource｜>'], + system_prefix=['<｜hy_begin_of_sentence:opensource｜>{{SYSTEM}}'], + prompt=['<｜hy_User:opensource｜>{{QUERY}}<｜hy_Assistant:opensource｜>'], + chat_sep=['<｜hy_eos:opensource｜>'], + suffix=['<｜hy_eos:opensource｜>'], + template_cls=HyV3Template, + is_thinking=True, + thinking_prefix='', + non_thinking_prefix='', + history_thinking_prefix='', + agent_template='hy_v3')) + + +class GptTemplate(Template): + support_padding_free = False + + def _get_gpt_oss_prefix(self): + today = datetime.now().strftime('%Y-%m-%d') + return ('<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\n' + f'Knowledge cutoff: 2024-06\nCurrent date: {today}\n\nReasoning: medium\n\n' + '# Valid channels: analysis, commentary, final. ' + 'Channel must be included for every message.<|end|>') + + def _swift_prepare_inputs(self, inputs: StdTemplateInputs): + super()._swift_prepare_inputs(inputs) + messages = inputs.messages + if self.use_chat_template: + if inputs.system is None: + inputs.system = self._get_gpt_oss_prefix() + elif not inputs.system.startswith('<|start|>'): + inputs.system = self._get_gpt_oss_prefix() + ( + f'<|start|>developer<|message|># Instructions\n\n{inputs.system}<|end|>') + for i, message in enumerate(messages): + if message['role'] == 'assistant' and isinstance(message['content'], str): + if not message['content'].startswith('<|channel|>'): + message['content'] = '<|channel|>final<|message|>' + message['content'] + + +@dataclass +class GptOssTemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=lambda: ['{{SYSTEM}}']) + prompt: Prompt = field(default_factory=lambda: ['<|start|>user<|message|>{{QUERY}}<|end|><|start|>assistant']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|end|>']) + suffix: Prompt = field(default_factory=lambda: ['<|return|>']) + + +register_template(GptOssTemplateMeta(LLMTemplateType.gpt_oss, template_cls=GptTemplate)) + +register_template( + TemplateMeta( + LLMTemplateType.longchat, + prefix=[], + system_prefix=['SYSTEM:{{SYSTEM}}'], + prompt=[' [Round {{ROUND0}}] USER:{{QUERY}} ASSISTANT:'], + chat_sep=[''], + suffix=[''], + )) + +register_template( + TemplateMeta( + LLMTemplateType.ling2, + prefix=['SYSTEMdetailed thinking off<|role_end|>'], + system_prefix=['SYSTEM{{SYSTEM}}\ndetailed thinking off<|role_end|>'], + prompt=['HUMAN{{QUERY}}<|role_end|>ASSISTANT'], + chat_sep=['<|role_end|>'], + suffix=['<|role_end|>'], + )) + +register_template( + TemplateMeta( + LLMTemplateType.ring2, + prefix=[], + system_prefix=['SYSTEM{{SYSTEM}}'], + prompt=['HUMAN{{QUERY}}ASSISTANT'], + chat_sep=[], + suffix=['<|endoftext|>'], + is_thinking=True, + thinking_prefix='\n', + )) + +register_template( + TemplateMeta( + LLMTemplateType.ring2_5, + prefix=[], + system_prefix=['SYSTEM\n{{SYSTEM}}\n\n'], + prompt=['HUMAN\n{{QUERY}}<|role_end|>\n\nASSISTANT\n'], + chat_sep=['<|role_end|>\n\n'], + suffix=['<|role_end|>\n\n'], + is_thinking=True, + )) + +register_template( + QwenTemplateMeta( + LLMTemplateType.iquestcoder, + default_system='You are LoopCoder, a helpful assistant developed by IQuest.', + )) + + +class YoutuLLMTemplate(Template): + + def _remove_thinking_content(self, content: str) -> str: + if '' in content: + content = content.rsplit('', 1)[-1].lstrip('\n') + return self.template_meta.history_thinking_prefix + content.strip() + + def _add_non_thinking_prefix(self, inputs) -> None: + messages = inputs.messages + non_thinking_prefix = self.template_meta.non_thinking_prefix + if non_thinking_prefix and messages: + # Find the last assistant message + for i in range(len(messages) - 1, -1, -1): + message = messages[i] + if message['role'] == 'assistant' and isinstance(message['content'], str): + if '' not in message['content'] and '' not in message['content']: + message['content'] = non_thinking_prefix + message['content'] + break + + def _remove_history_thinking(self, inputs) -> None: + messages = inputs.messages + first_tool_index = len(messages) + for i, message in enumerate(messages): + if message['role'] == 'tool' or (message['role'] == 'user' and isinstance(message.get('content'), str) + and message['content'].startswith('') + and message['content'].endswith('')): + first_tool_index = i + break + # Only remove thinking content for assistant messages before first_tool_index - 1 + for i, message in enumerate(messages): + if message['role'] == 'assistant' and isinstance(message['content'], str): + is_last = (i == len(messages) - 1) + if not is_last and i < first_tool_index - 1: + message['content'] = self._remove_thinking_content(message['content']) + + +register_template( + TemplateMeta( + LLMTemplateType.youtu_llm, + template_cls=YoutuLLMTemplate, + prefix=[['bos_token_id']], + system_prefix=[['bos_token_id'], '{{SYSTEM}}'], + prompt=['<|User|>{{QUERY}}<|Assistant|>'], + chat_sep=['<|end_of_text|>'], + suffix=['<|end_of_text|>'], + is_thinking=True, + non_thinking_prefix='\n\n\n\n', + agent_template='youtu', + )) + +register_template( + TemplateMeta( + LLMTemplateType.olmoe, + prefix=['|||IP_ADDRESS|||'], + system_prefix=['|||IP_ADDRESS|||<|system|>\n{{SYSTEM}}\n'], + prompt=['<|user|>\n{{QUERY}}\n<|assistant|>\n'], + chat_sep=['|||IP_ADDRESS|||\n'], + suffix=['|||IP_ADDRESS|||'], + stop_words=['<|endoftext|>'], + )) + +register_template( + TemplateMeta( + LLMTemplateType.olmoe_0924, + prefix=['<|endoftext|>'], + system_prefix=['<|endoftext|><|system|>\n{{SYSTEM}}\n'], + prompt=['<|user|>\n{{QUERY}}\n<|assistant|>\n'], + chat_sep=['<|endoftext|>\n'], + suffix=['<|endoftext|>'], + stop_words=['<|endoftext|>'], + )) + +register_template( + QwenTemplateMeta( + LLMTemplateType.nemotron_h, + default_system=None, + is_thinking=True, + thinking_prefix='\n', + stop_words=['<|im_end|>'], + agent_template='hermes', + )) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/megrez.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/megrez.py new file mode 100644 index 0000000000000000000000000000000000000000..5c8748c53641468a25ae4ed80af25ea5ebd052a6 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/megrez.py @@ -0,0 +1,96 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +import torch.nn as nn +from dataclasses import dataclass, field +from typing import Any, Dict, List, Literal, Optional + +from ..base import Template +from ..constant import LLMTemplateType, MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, findall + + +@dataclass +class MegrezTemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=lambda: ['<|role_start|>system<|role_end|>{{SYSTEM}}<|turn_end|>']) + prompt: Prompt = field(default_factory=lambda: + ['<|role_start|>user<|role_end|>{{QUERY}}<|turn_end|><|role_start|>assistant<|role_end|>']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|turn_end|>']) + suffix: Prompt = field(default_factory=lambda: ['<|turn_end|>']) + default_system: str = '你是Megrez-3B-Instruct，将针对用户的问题给出详细的、积极的回答。' + + +register_template(MegrezTemplateMeta(LLMTemplateType.megrez)) + + +class MegrezOmniTemplate(Template): + skip_prompt = False + placeholder_tokens = ['<|unk|>'] + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + return [[-1], '\n'] + elif media_type == 'audio': + return [f'Audio {index + 1}: ', [-2], '\n'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + + for mm_key in ['images', 'audios']: + mm_data = getattr(inputs, mm_key) + if not mm_data: + continue + if mm_key == 'images': + idx_list = findall(input_ids, -1) + encoding = self.processor.process_image( + mm_data, + return_tensors='pt', + ) + text = self.processor.insert_image_feature_placeholders( + ''.join(['(./)'] * len(mm_data)), encoding) + encoded['image_encoding'] = encoding + else: + idx_list = findall(input_ids, -2) + encoding = self.processor.process_audio( + mm_data, + return_tensors='pt', + ) + text = self.processor.insert_audio_feature_placeholders( + ''.join(['()'] * len(mm_data)), encoding) + encoded['audio_encoding'] = encoding + + padding = text.split('') + + def _get_new_tokens(i): + return self._tokenize(padding[i]) + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + return encoded + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + _, inputs_embeds, _ = model.compose_embeddings(inputs) + inputs.pop('position_ids', None) + return {'inputs_embeds': inputs_embeds} + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + new_batch = [] + for b in batch: + text_encodings = {'input_ids': torch.tensor(b['input_ids'])} + multimodal_inputs = {'image_encoding': b.get('image_encoding'), 'audio_encoding': b.get('audio_encoding')} + new_batch.append(self.processor.merge_encodings(text_encodings, multimodal_inputs)) + res.update(self.processor.data_collator(new_batch)) + return res + + +register_template(MegrezTemplateMeta(MLLMTemplateType.megrez_omni, template_cls=MegrezOmniTemplate)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/microsoft.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/microsoft.py new file mode 100644 index 0000000000000000000000000000000000000000..8cd16d92daf773749863fb66ba3b07ae1af398e4 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/microsoft.py @@ -0,0 +1,209 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import json +import torch +from dataclasses import dataclass, field +from torch import nn +from typing import Any, Dict, List, Literal, Optional + +from ..base import Template +from ..constant import LLMTemplateType, MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, findall +from ..vision_utils import load_file + + +class FlorenceTemplate(Template): + # If it's an encoder-decoder architecture, the default settings are + # loss_scale: 'last_round' and skip_prompt: False. + is_encoder_decoder = True + skip_prompt = False + + @staticmethod + def _add_default_tags(inputs: StdTemplateInputs) -> None: + return + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + return [] + + def replace_bbox(self, bbox: List[int], index: int, inputs: StdTemplateInputs) -> List[Context]: + return [''.join(f'' for box in bbox)] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + processor = self.processor + inputs.query = inputs.to_history()['query'] + new_query = processor._construct_prompts([inputs.query])[0] + for i in reversed(range(len(inputs.messages))): + if inputs.messages[i]['role'] == 'user': + inputs.messages[i]['content'] = new_query + break + encoded = super()._encode(inputs) + input_ids = encoded['prompt_input_ids'] + images = inputs.images or [] + labels = encoded['answer_labels'] + if labels is not None: + labels = [0] + labels + if images: + pixel_values = processor.image_processor( + images, return_tensors='pt')['pixel_values'].to(self.model_info.torch_dtype) + encoded['pixel_values'] = pixel_values + encoded['input_ids'] = input_ids + encoded['labels'] = labels + return encoded + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + inputs_embeds = model.get_input_embeddings()(inputs['input_ids']) + pixel_values = inputs.get('pixel_values') + if pixel_values is not None: + image_features = model._encode_image(pixel_values) + inputs_embeds, inputs['attention_mask'] = model._merge_input_ids_with_image_features( + image_features, inputs_embeds) + return {'inputs_embeds': inputs_embeds} + + def decode_generate_ids(self, generate_ids: List[int], **kwargs) -> Any: + response = super().decode_generate_ids(generate_ids, **kwargs) + template_inputs = kwargs.get('template_inputs') + images = template_inputs.images + image_size = None + if images: + image_size = (images[0].width, images[0].height) + query_before, query_sep, query_after = template_inputs.query.partition('>') + task = query_before + query_sep if query_sep else '' + return json.dumps(self.processor.post_process_generation(response, task=task, image_size=image_size)) + + +register_template( + TemplateMeta( + MLLMTemplateType.florence, + prefix=[''], + prompt=['{{QUERY}}'], + chat_sep=None, + suffix=[''], + template_cls=FlorenceTemplate, + )) + + +@dataclass +class Phi3TemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=list) + prompt: Prompt = field(default_factory=lambda: ['<|user|>\n{{QUERY}}<|end|>\n<|assistant|>\n']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|end|>\n']) + suffix: Prompt = field(default_factory=lambda: ['<|end|>']) + system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<|system|>\n{{SYSTEM}}<|end|>\n']) + auto_add_bos: bool = True + + +register_template(Phi3TemplateMeta(LLMTemplateType.phi3)) + + +@dataclass +class Phi4TemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=list) + prompt: Prompt = field( + default_factory=lambda: ['<|im_start|>user<|im_sep|>{{QUERY}}<|im_end|><|im_start|>assistant<|im_sep|>']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|im_end|>']) + suffix: Prompt = field(default_factory=lambda: ['<|im_end|>']) + system_prefix: Optional[Prompt] = field( + default_factory=lambda: ['<|im_start|>system<|im_sep|>{{SYSTEM}}<|im_end|>']) + auto_add_bos: bool = True + + +register_template(Phi4TemplateMeta(LLMTemplateType.phi4)) + + +class Phi3VisionTemplate(Template): + image_placeholder = ['<|image|>\n'] # <|image|>\n + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if self.mode == 'vllm': + return [f'<|image_{index + 1}|>\n'] # <|image_1|>\n + else: + return super().replace_tag(media_type, index, inputs) + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + images = inputs.images or [] + encoded = super()._encode(inputs) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + idx_list = findall(input_ids, 32044) # '<|image|>' + + if len(images) > 0: + processor = self.processor + encoded.update(processor.image_processor(images, return_tensors='pt')) + assert len(idx_list) == len(images), f'len(idx_list): {len(idx_list)}, len(images): {len(images)}' + res_input_ids = [] + res_labels = [] + num_img_tokens = encoded.pop('num_img_tokens').tolist() + idx_list.insert(0, -1) + for i in range(len(idx_list) - 1): + image_token_id = -i - 1 + res_input_ids += input_ids[idx_list[i] + 1:idx_list[i + 1]] + [image_token_id] * num_img_tokens[i] + if labels is not None: + res_labels += labels[idx_list[i] + 1:idx_list[i + 1]] + [-100] * num_img_tokens[i] + res_input_ids += input_ids[idx_list[-1] + 1:] + input_ids = res_input_ids + if labels is not None: + res_labels += labels[idx_list[-1] + 1:] + labels = res_labels + + encoded['input_ids'] = input_ids + encoded['labels'] = labels + return encoded + + +class Phi4MMTemplate(Template): + placeholder_tokens = ['<|endoftext10|>', '<|endoftext11|>'] + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + if self.mode == 'vllm': + return [f'<|image_{index + 1}|>'] # <|image_1|> + return [[-100]] + elif media_type == 'audio': + import soundfile as sf + inputs.audios[index] = sf.read(load_file(inputs.audios[index])) + return [[-200]] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + images_idx = findall(input_ids, -100) + audios_idx = findall(input_ids, -200) + text = '\n'.join(['<|image_1|>'] * len(inputs.images) + ['<|audio_1|>'] * len(inputs.audios)) + new_encoded = self.processor( + text=text, images=inputs.images or None, audios=inputs.audios or None, return_tensors='pt') + placeholders = self._split_list(new_encoded.pop('input_ids')[0].tolist(), 198) + + def _get_new_tokens(i): + return placeholders[i] + + encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens( + input_ids, labels, loss_scale, images_idx + audios_idx, _get_new_tokens) + new_encoded.pop('attention_mask') + encoded.update(new_encoded) + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + keys = [ + 'input_image_embeds', 'image_sizes', 'image_attention_mask', 'input_audio_embeds', 'audio_embed_sizes', + 'input_mode' + ] + inputs = self.fetch_inputs(batch, keys) + for k, v in inputs.items(): + inputs[k] = torch.concat(v) + res.update(inputs) + return res + + +register_template(Phi3TemplateMeta(MLLMTemplateType.phi3_vision, template_cls=Phi3VisionTemplate)) + +register_template(Phi3TemplateMeta( + MLLMTemplateType.phi4_multimodal, + template_cls=Phi4MMTemplate, +)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/midashenglm.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/midashenglm.py new file mode 100644 index 0000000000000000000000000000000000000000..078a5fbc7f5e01057d000d1e36b32fdedf582581 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/midashenglm.py @@ -0,0 +1,65 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +import torch.nn.functional as F +from functools import partial +from typing import Any, Dict, List, Literal, Optional + +from swift.utils import get_env_args +from ..base import Template +from ..constant import MLLMTemplateType +from ..register import register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, findall +from ..vision_utils import load_batch +from .qwen import QwenTemplateMeta + + +class MiDashengLMTemplate(Template): + placeholder_tokens = ['<|AUDIO|>'] + skip_prompt = False + + def init_env_args(self): + super().init_env_args() + self.sampling_rate = get_env_args('sampling_rate', int, 16000) + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'audio' + return ['<|AUDIO|>'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + from transformers.audio_utils import load_audio + encoded = super()._encode(inputs) + input_ids = encoded['input_ids'] + inputs.audios = load_batch(inputs.audios, partial(load_audio, sampling_rate=self.sampling_rate)) + audio_token = self._tokenize('<|AUDIO|>')[0] + idx_list = findall(input_ids, audio_token) + if idx_list: + split_token = self._tokenize('\n')[0] + audio_inputs = self.processor(text='\n'.join(['<|AUDIO|>'] * len(inputs.audios)), audio=inputs.audios) + splited_tokens = self._split_list(audio_inputs['input_ids'][0].tolist(), split_token) + + encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens( + input_ids, encoded['labels'], encoded['loss_scale'], idx_list, lambda i: splited_tokens[i]) + encoded['input_values'] = audio_inputs['input_values'] + encoded['audio_length'] = audio_inputs['audio_length'] + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + + input_values = [b['input_values'] for b in batch if b.get('input_values') is not None] + audio_lengths = [b['audio_length'] for b in batch if b.get('audio_length') is not None] + + if input_values: + res['audio_length'] = torch.concat(audio_lengths) + for i in range(len(input_values)): + # TODO: check padding_side + pad_len = (res['audio_length'].max() - input_values[i].shape[1]).item() + input_values[i] = F.pad(input_values[i], (0, pad_len), 'constant', 0) + res['input_values'] = torch.concat(input_values) + + return res + + +register_template(QwenTemplateMeta(MLLMTemplateType.midashenglm, template_cls=MiDashengLMTemplate)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mimo.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mimo.py new file mode 100644 index 0000000000000000000000000000000000000000..846d58199439c184f63fff55eddb467063bceae6 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mimo.py @@ -0,0 +1,208 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import sys +import torch +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from ..base import Template +from ..constant import MLLMTemplateType +from ..register import register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, Word, findall +from ..vision_utils import load_audio +from .utils import ChatmlTemplateMeta + + +@dataclass +class MiMoV2TemplateMeta(ChatmlTemplateMeta): + default_system: Optional[str] = 'You are MiMo, a helpful AI assistant engineered by Xiaomi.' + auto_add_bos: bool = False + stop_words: List[Word] = field(default_factory=lambda: ['<|endoftext|>']) + prefix: Prompt = field(default_factory=list) + prompt: Prompt = field(default_factory=lambda: ['<|im_start|>user\n{{QUERY}}<|im_end|><|im_start|>assistant\n']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|im_end|>']) + suffix: Prompt = field(default_factory=lambda: ['<|im_end|>']) + system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<|im_start|>system\n{{SYSTEM}}<|im_end|>']) + + +def _wav_to_log_mel(wav: torch.Tensor, *, sr, n_mels, n_fft, hop, win, f_min, f_max) -> torch.Tensor: + from torchaudio.transforms import MelSpectrogram + + mel_transform = MelSpectrogram( + sample_rate=sr, + n_fft=n_fft, + hop_length=hop, + win_length=win, + f_min=f_min, + f_max=f_max, + n_mels=n_mels, + power=1.0, + center=True, + ) + spec = mel_transform(wav[None, :]) + return torch.log(torch.clip(spec, min=1e-7)).squeeze().transpose(0, 1) + + +def _pad_codes_to_group(codes: torch.Tensor, group_size: int, audio_channels: int) -> torch.Tensor: + codes = codes[:, :audio_channels] + T = codes.shape[0] + pad_T = ((T + group_size - 1) // group_size) * group_size + if pad_T > T: + codes = torch.cat([codes, codes[-1:].expand(pad_T - T, -1)], dim=0) + return codes + + +class MiMoV2Template(Template): + image_token_id = 151655 + video_token_id = 151656 + audio_token_id = 151669 + placeholder_tokens = ['<|image_pad|>', '<|video_pad|>', '<|audio_pad|>'] + norm_bbox = 'none' + use_model = True + + def init_env_args(self) -> None: + super().init_env_args() + pc = getattr(self.config, 'processor_config', None) or {} + self.audio_sampling_rate = pc.get('audio_sampling_rate', 24000) + + @property + def _audio_cfg(self) -> dict: + return getattr(self.config, 'audio_config', None) or {} + + def replace_tag(self, media_type, index, inputs: StdTemplateInputs) -> List[Context]: + from qwen_vl_utils import fetch_image, fetch_video + + assert media_type in {'image', 'video', 'audio'} + if media_type == 'audio': + return ['<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>'] + + kwargs = {'image_patch_size': self.processor.image_processor.patch_size} + if media_type == 'image': + inputs.images[index] = fetch_image({'image': inputs.images[index], **inputs.chat_template_kwargs}, **kwargs) + if self.mode == 'lmdeploy': + return ['<|vision_start|>', [-100], '<|vision_end|>'] + return ['<|vision_start|><|image_pad|><|vision_end|>'] + else: + if self.mode == 'sglang': + return ['<|vision_start|><|video_pad|><|vision_end|>'] + video = inputs.videos[index] + video_inputs = {'video': video, **inputs.chat_template_kwargs} + if isinstance(video, list): + from qwen_vl_utils import vision_process + video_inputs['sample_fps'] = vision_process.FPS + video, _ = fetch_video(video_inputs, return_video_sample_fps=True) + if isinstance(video, torch.Tensor): + video = video.to(torch.uint8) + inputs.videos[index] = video + return ['<|vision_start|><|video_pad|><|vision_end|>'] + + def _encode_truncated(self, inputs: StdTemplateInputs): + encoded = super()._encode_truncated(inputs) + if self.mode == 'sglang': + batched = encoded if isinstance(encoded, list) else [encoded] + for item in batched: + for old, new in [('images', 'image_data'), ('audios', 'audio_data'), ('videos', 'video_data')]: + if old in item: + item[new] = item.pop(old) + for key in ['labels', 'loss_scale', 'channel']: + item.pop(key, None) + mm_keys = {'image_data', 'audio_data', 'video_data'} + if item.keys() & mm_keys and 'input_ids' in item: + ids = item['input_ids'] + if hasattr(ids, 'tolist'): + ids = ids.tolist() + item['prompt'] = self.tokenizer.decode(ids, skip_special_tokens=False) + return encoded + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + processor = self.processor + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + + if inputs.audios: + audio_tokenizer = self.model.audio_tokenizer + at_config = audio_tokenizer.config + group_size = self._audio_cfg.get('group_size', 4) + audio_channels = self._audio_cfg.get('audio_channels', 20) + segment_size = self._audio_cfg.get('audio_segment_size', 6000) + + # wav → mel + mels = [ + _wav_to_log_mel( + torch.from_numpy(load_audio(p, self.audio_sampling_rate)).float(), + sr=at_config.sampling_rate, + n_mels=at_config.n_mels, + n_fft=at_config.nfft, + hop=at_config.hop_length, + win=at_config.window_size, + f_min=at_config.fmin, + f_max=at_config.fmax, + ) for p in inputs.audios + ] + + # mel → codes + tokenize_audio_batch = sys.modules[type(self.model).__module__].tokenize_audio_batch + code_list = tokenize_audio_batch( + mels, + audio_tokenizer.encoder, + segment_size=segment_size, + device=audio_tokenizer.device, + ) + + # pad per-audio so placeholder count matches model's per-group output + padded_codes_list = [_pad_codes_to_group(c, group_size, audio_channels) for c in code_list] + + idx_list = findall(input_ids, self.audio_token_id) + + def _get_new_audio_tokens(i): + return [self.audio_token_id] * (padded_codes_list[i].shape[0] // group_size) + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_audio_tokens) + encoded['audio_codes'] = torch.cat(padded_codes_list, dim=0) + + for media_type in ['images', 'videos']: + mm_data = getattr(inputs, media_type) + if not mm_data: + continue + if media_type == 'images': + media_token = self.image_token_id + media_inputs = processor.image_processor(images=mm_data, return_tensors='pt', do_resize=False) + media_grid_thw = media_inputs['image_grid_thw'] + else: + if hasattr(processor, 'video_processor'): + processor_func = processor.video_processor + else: + processor_func = processor.image_processor + media_inputs = processor_func(videos=mm_data, return_tensors='pt', do_resize=False) + media_grid_thw = media_inputs['video_grid_thw'] + media_token = self.video_token_id + + idx_list = findall(input_ids, media_token) + merge_length = processor.image_processor.merge_size**2 + + def _get_new_tokens(i): + return [media_token] * (media_grid_thw[i].prod() // merge_length) + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + encoded.update(media_inputs) + + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + if 'pixel_values_videos' in res: + res['video_pixel_values'] = res.pop('pixel_values_videos') + audio_codes = [b['audio_codes'] for b in batch if b.get('audio_codes') is not None] + if audio_codes: + res['audio_codes'] = torch.cat(audio_codes, dim=0) + return res + + +register_template(MiMoV2TemplateMeta(MLLMTemplateType.mimo_v2, template_cls=MiMoV2Template)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minicpm.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minicpm.py new file mode 100644 index 0000000000000000000000000000000000000000..19cc64ea78bcec568d9abb74a5d8a90adc41757b --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minicpm.py @@ -0,0 +1,723 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import math +import numpy as np +import torch +import transformers +from dataclasses import dataclass, field +from functools import partial +from packaging import version +from torch import nn +from typing import Any, Dict, List, Literal, Optional + +from swift.utils import get_env_args, get_packed_seq_params +from ..base import Template +from ..constant import LLMTemplateType, MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, findall +from ..vision_utils import load_audio, load_video_minicpmv_mplug_owl3 +from .llama import Llama3TemplateMeta +from .qwen import Qwen2_5TemplateMeta, Qwen3MixedTemplateMeta, QwenTemplateMeta +from .utils import ChatmlTemplateMeta + + +@dataclass +class MinicpmTemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=lambda: ['{{SYSTEM}}']) + prompt: Prompt = field(default_factory=lambda: ['<用户>{{QUERY}}']) + chat_sep: Optional[Prompt] = field(default_factory=list) + suffix: Prompt = field(default_factory=lambda: ['']) + + +register_template(MinicpmTemplateMeta(LLMTemplateType.minicpm)) + + +def _remove_idx(arr: List[int], idx_list: List[int]) -> List[int]: + res = [] + idx_set = set(idx_list) + for i, x in enumerate(arr): + if i not in idx_set: + res.append(x) + return res + + +class MiniCPMVTemplate(Template): + is_v2_5 = False + use_model = True + skip_prompt = False + placeholder_tokens = [''] + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if self.mode == 'vllm': + return ['(./)\n'] + else: + return [[-100]] + + async def prepare_lmdeploy_turbomind_inputs(self, inputs: Dict[str, Any]) -> None: + images = inputs.pop('images', None) or [] + if len(images) == 0: + return + input_ids = inputs['input_ids'] + idx_list = findall(input_ids, -100) + idx_list.insert(0, -1) + new_input_ids = [] + features = [] + for i in range(len(idx_list) - 1): + new_input_ids += input_ids[idx_list[i] + 1:idx_list[i + 1]] + context_list = ['', [-100], ''] + feat = [x.squeeze() for x in images[i]['embeddings'].split(1)] + grid = images[i].get('grid') + if len(feat) > 1 and grid is not None: + context_list.append('') + for j in range(grid[1]): + if j > 0: + context_list.append('\n') + for _ in range(grid[0]): + context_list += ['', [-100], ''] + context_list.append('\n') + new_input_ids += self._encode_context_list(context_list)[0] + features += feat + new_input_ids += input_ids[idx_list[-1] + 1:] + inputs['input_ids'] = new_input_ids + inputs['images'] = features + await super().prepare_lmdeploy_turbomind_inputs(inputs) + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + images = inputs.images + input_ids = encoded['input_ids'] + labels = encoded['labels'] + idx_list = findall(input_ids, -100) + idx = idx_list[0] + tgt_sizes = None + slice_mode = getattr(self.config, 'slice_mode', False) + if slice_mode: + if self.is_v2_5: + image_processor = self.processor.image_processor + image_inputs = image_processor(images, return_tensors='pt').to(self.model_info.torch_dtype) + placeholder = image_processor.get_slice_image_placeholder(image_inputs.image_sizes[0][0]) + pixel_values = image_inputs['pixel_values'] + tgt_sizes = image_inputs['tgt_sizes'] + else: + images, placeholder = self.model.get_slice_image_placeholder(images[0], self.processor) + pixel_values = [[self.model.transform(img) for img in images]] + placeholder += '\n' + placeholder_id = self.processor.encode(placeholder, add_special_tokens=False) + input_ids = (input_ids[:idx] + placeholder_id + input_ids[idx + 1:]) + if labels is not None: + labels = (labels[:idx] + [-100] * len(placeholder_id) + labels[idx + 1:]) + input_tensor_ids = torch.tensor(input_ids) + image_start_idx = torch.where(input_tensor_ids == self.processor.im_start_id)[0] + image_start_idx += 1 + image_end_idx = torch.where(input_tensor_ids == self.processor.im_end_id)[0] + valid_image_nums = max(len(image_start_idx), len(image_end_idx)) + image_bound = [ + torch.hstack( + [image_start_idx[:valid_image_nums].unsqueeze(-1), image_end_idx[:valid_image_nums].unsqueeze(-1)]) + ] + else: + placeholder = '' + '' * self.config.query_num + '\n' + placeholder_id = self.processor.encode(placeholder, add_special_tokens=False) + input_ids = (input_ids[:idx] + placeholder_id + input_ids[idx + 1:]) + if labels is not None: + labels = (labels[:idx] + [-100] * len(placeholder_id) + labels[idx + 1:]) + image_bound = [torch.tensor([[idx, idx + self.config.query_num]])] + pixel_values = [[self.model.transform(images[0])]] + encoded = { + 'input_ids': input_ids, + 'labels': labels, + 'image_bound': image_bound, + 'pixel_values': pixel_values, + 'tgt_sizes': tgt_sizes + } + return encoded + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + inputs_embeds, _ = model.get_vllm_embedding(inputs) + return {'inputs_embeds': inputs_embeds} + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = {} + for k in ['pixel_values', 'image_bound', 'tgt_sizes']: + res[k] = self.gather_list(batch, k) + res.update(super()._data_collator(batch, padding_to=padding_to)) + return res + + +register_template(MinicpmTemplateMeta(MLLMTemplateType.minicpmv, template_cls=MiniCPMVTemplate)) + + +class MiniCPMV2_5Template(MiniCPMVTemplate): + is_v2_5 = True + + +register_template(Llama3TemplateMeta( + MLLMTemplateType.minicpmv2_5, + template_cls=MiniCPMV2_5Template, +)) + + +class MiniCPMV2_6Template(MiniCPMVTemplate): + + def init_env_args(self): + super().init_env_args() + self.max_num_frames = get_env_args('max_num_frames', int, 64) + self.max_slice_nums = get_env_args('max_slice_nums', int, None) + self.video_max_slice_nums = get_env_args('video_max_slice_nums', int, 1) # or 2 + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type in {'image', 'video'} + load_video = partial(load_video_minicpmv_mplug_owl3, max_num_frames=self.max_num_frames) + image_context = super().replace_tag('image', index, inputs) + if media_type == 'image': + return image_context + elif media_type == 'video': + return self.replace_video2image(load_video, inputs, lambda i: image_context) + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = Template._encode(self, inputs) + images = inputs.images + use_video = bool(inputs.videos) + use_image_id = True + max_slice_nums = self.max_slice_nums + if use_video: + max_slice_nums = self.video_max_slice_nums + use_image_id = False + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + idx_list = findall(input_ids, -100) + + image_processor = self.processor.image_processor + image_inputs = image_processor([images], return_tensors='pt', + max_slice_nums=max_slice_nums).to(self.model_info.torch_dtype) + + def _get_new_tokens(i): + placeholder = image_processor.get_slice_image_placeholder( + image_inputs.image_sizes[0][i], image_idx=i, max_slice_nums=max_slice_nums, use_image_id=use_image_id) + placeholder += '\n' + return self.processor.encode(placeholder, add_special_tokens=False) + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, _get_new_tokens) + + if inputs.images: + input_tensor_ids = torch.tensor(input_ids) + unk_token = self.processor.encode('', add_special_tokens=False)[0] + indices = (input_tensor_ids == unk_token).nonzero(as_tuple=True)[0].tolist() + + ranges = [] + start = indices[0] + for i in range(1, len(indices)): + if indices[i] != indices[i - 1] + 1: + ranges.append([start, indices[i - 1] + 1]) + start = indices[i] + ranges.append([start, indices[-1] + 1]) + image_bound = [torch.tensor(ranges)] + else: + image_bound = [[]] + + encoded = { + 'input_ids': input_ids, + 'labels': labels, + 'loss_scale': loss_scale, + 'image_bound': image_bound, + 'pixel_values': image_inputs['pixel_values'], + 'tgt_sizes': image_inputs['tgt_sizes'] + } + return encoded + + +register_template(QwenTemplateMeta( + MLLMTemplateType.minicpmv2_6, + template_cls=MiniCPMV2_6Template, +)) + +register_template(ChatmlTemplateMeta( + MLLMTemplateType.minicpmv4, + template_cls=MiniCPMV2_6Template, +)) + +register_template(Qwen2_5TemplateMeta( + MLLMTemplateType.minicpmo, + template_cls=MiniCPMV2_6Template, +)) + + +class MiniCPMV4_5Template(MiniCPMV2_6Template): + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = Template._encode(self, inputs) + images = inputs.images + use_video = bool(inputs.videos) + use_image_id = True + max_slice_nums = self.max_slice_nums + if use_video: + max_slice_nums = self.video_max_slice_nums + use_image_id = False + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + idx_list = findall(input_ids, -100) + + image_processor = self.processor.image_processor + image_inputs = image_processor([images], return_tensors='pt', + max_slice_nums=max_slice_nums).to(self.model_info.torch_dtype) + + def _get_new_tokens(i): + placeholder = image_processor.get_slice_image_placeholder( + image_inputs.image_sizes[0][i], image_idx=i, max_slice_nums=max_slice_nums, use_image_id=use_image_id) + placeholder += '\n' + return self.processor.encode(placeholder, add_special_tokens=False) + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, _get_new_tokens) + + if inputs.images: + input_tensor_ids = torch.tensor(input_ids) + unk_token = self.processor.encode('', add_special_tokens=False)[0] + indices = (input_tensor_ids == unk_token).nonzero(as_tuple=True)[0].tolist() + + ranges = [] + start = indices[0] + for i in range(1, len(indices)): + if indices[i] != indices[i - 1] + 1: + ranges.append([start, indices[i - 1] + 1]) + start = indices[i] + ranges.append([start, indices[-1] + 1]) + image_bound = [torch.tensor(ranges)] + else: + image_bound = [[]] + + encoded = { + 'input_ids': input_ids, + 'labels': labels, + 'loss_scale': loss_scale, + 'image_bound': image_bound, + 'pixel_values': image_inputs['pixel_values'], + 'tgt_sizes': image_inputs['tgt_sizes'], + 'temporal_ids': image_inputs['temporal_ids'], + } + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = {} + for k in ['pixel_values', 'image_bound', 'tgt_sizes', 'temporal_ids']: + res[k] = self.gather_list(batch, k) + res.update(Template._data_collator(self, batch, padding_to=padding_to)) + return res + + +register_template( + Qwen3MixedTemplateMeta( + MLLMTemplateType.minicpmv4_5, + template_cls=MiniCPMV4_5Template, + is_thinking=True, + thinking_prefix='\n', + )) + + +class MiniCPMO4_5Template(MiniCPMV4_5Template): + SAMPLING_RATE = 16000 + MAX_AUDIO_DURATION = 30 # seconds + + def init_env_args(self): + super().init_env_args() + self.use_audio_in_video = get_env_args('use_audio_in_video', bool, False) + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image' or media_type == 'video' and not self.use_audio_in_video: + return super().replace_tag(media_type, index, inputs) + elif media_type == 'audio': + # Load audio from file path to numpy array at 16kHz + if isinstance(inputs.audios[index], str): + inputs.audios[index] = load_audio(inputs.audios[index], sampling_rate=self.SAMPLING_RATE) + return ['<|audio_start|><|audio_end|>'] + elif media_type == 'video': + from minicpmo.utils import get_video_frame_audio_segments + video = inputs.videos[inputs.video_idx] + video_segments, audio_segments, _ = get_video_frame_audio_segments( + video, use_audio=self.use_audio_in_video, stack_frames=1) + # Insert frames into images list at current position + images = inputs.images + inputs.images = images[:inputs.image_idx] + video_segments + images[inputs.image_idx:] + # Build context list + image_context = [[-100]] + context_list = [] + if self.use_audio_in_video and audio_segments: + # Insert audio segments into audios list at current position + audios = inputs.audios + inputs.audios = audios[:inputs.audio_idx] + audio_segments + audios[inputs.audio_idx:] + audio_context = ['<|audio_start|><|audio_end|>'] + # Interleave: one image placeholder + one audio placeholder per second + for i in range(len(video_segments)): + context_list += image_context + if i < len(audio_segments): + context_list += audio_context + inputs.audio_idx += len(audio_segments) + else: + for _ in range(len(video_segments)): + context_list += image_context + inputs.image_idx += len(video_segments) + return context_list + + def _get_audio_num_tokens(self, audio_sample_len: int) -> int: + """Compute the number of placeholder tokens for an audio of given sample count. + + This mirrors the official get_audio_placeholder logic: + 1. mel frames = ceil(audio_samples / hop_length) + 2. after CNN downsampling: (mel_frames - 1) // 2 + 1 + 3. after avg pooling: (cnn_frames - pool_step) // pool_step + 1 + """ + hop_length = self.processor.audio_processor.hop_length # 160 + pool_step = self.config.audio_pool_step # 5 + feature_lens = math.ceil(audio_sample_len / hop_length) + feature_lens_after_cnn = (feature_lens - 1) // 2 + 1 + output_lens = (feature_lens_after_cnn - pool_step) // pool_step + 1 + return output_lens + + def _extract_audio_features(self, audios: List[np.ndarray]): + """Extract mel features from audio arrays using the WhisperFeatureExtractor. + + Handles chunking of long audios (>30s) into segments. + Matches the official audio_feature_extract output format. + + Returns: + audio_features: tensor [N, 80, max_frames] or [] if no audios + audio_feature_lens: [tensor([l1, l2, ...])] or None + """ + audio_processor = self.processor.audio_processor + max_audio_inp_len = self.MAX_AUDIO_DURATION * self.SAMPLING_RATE + all_audio_features = [] + all_audio_lens = [] + + for audio in audios: + # Chunk long audios at 30s boundaries + if len(audio) <= max_audio_inp_len: + chunks = [audio] + else: + chunks = [audio[i:i + max_audio_inp_len] for i in range(0, len(audio), max_audio_inp_len)] + + for chunk in chunks: + audio_input = audio_processor( + chunk, + sampling_rate=self.SAMPLING_RATE, + return_tensors='pt', + padding='max_length', + return_attention_mask=True, + ) + feat = audio_input['input_features'] # [1, 80, frames] + actual_len = audio_input['attention_mask'].sum(dim=1) # [1] + feat = feat[:, :, :actual_len[0]] + all_audio_features.append(feat.squeeze(0)) # [80, actual_frames] + all_audio_lens.append(actual_len[0]) + + if all_audio_features: + # Pad and stack: [N, 80, max_frames] — same as official processor + audio_features = torch.nn.utils.rnn.pad_sequence( + [f.transpose(0, 1) for f in all_audio_features], + batch_first=True, + padding_value=0.0, + ).transpose(1, 2) + audio_feature_lens = [torch.hstack(all_audio_lens)] + else: + audio_features = [] + audio_feature_lens = None + + return audio_features, audio_feature_lens + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + # Step 1: Base encode — produces input_ids with -100 for images + # and audio_start_id,audio_end_id pairs for audios + encoded = Template._encode(self, inputs) + images = inputs.images + use_video = bool(inputs.videos) + use_image_id = True + max_slice_nums = self.max_slice_nums + if use_video: + max_slice_nums = self.video_max_slice_nums + use_image_id = False + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + + # Step 2: Process images — replace -100 tokens with image placeholders + idx_list = findall(input_ids, -100) + image_processor = self.processor.image_processor + image_inputs = image_processor([images], return_tensors='pt', + max_slice_nums=max_slice_nums).to(self.model_info.torch_dtype) + + def _get_new_tokens(i): + placeholder = image_processor.get_slice_image_placeholder( + image_inputs.image_sizes[0][i], image_idx=i, max_slice_nums=max_slice_nums, use_image_id=use_image_id) + placeholder += '\n' + return self.processor.encode(placeholder, add_special_tokens=False) + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, _get_new_tokens) + + # Step 3: Process audios — expand audio_start/audio_end pairs with tokens + tokenizer = self.processor.tokenizer + audio_start_id = tokenizer.convert_tokens_to_ids('<|audio_start|>') + audio_end_id = tokenizer.convert_tokens_to_ids('<|audio_end|>') + unk_token_id = tokenizer.convert_tokens_to_ids('') + + audio_features = None + audio_feature_lens = None + + if inputs.audios: + audios = inputs.audios + audio_features, audio_feature_lens = self._extract_audio_features(audios) + + # Find empty audio placeholder pairs (audio_start_id immediately followed by audio_end_id) + audio_placeholder_positions = [] + for i in range(len(input_ids) - 1): + if input_ids[i] == audio_start_id and input_ids[i + 1] == audio_end_id: + audio_placeholder_positions.append(i) + + assert len(audio_placeholder_positions) == len(audios), \ + f'Found {len(audio_placeholder_positions)} audio placeholders but have {len(audios)} audios' + + # Expand each audio placeholder with tokens + offset = 0 + for i, audio in enumerate(audios): + num_tokens = self._get_audio_num_tokens(len(audio)) + unk_tokens = [unk_token_id] * num_tokens + pos = audio_placeholder_positions[i] + offset + # Current: [..., audio_start_id, audio_end_id, ...] + # Target: [..., audio_start_id, unk*N, audio_end_id, ...] + input_ids = input_ids[:pos + 1] + unk_tokens + input_ids[pos + 1:] + if labels is not None: + labels = labels[:pos + 1] + [-100] * num_tokens + labels[pos + 1:] + if loss_scale is not None: + scale_val = loss_scale[pos] + loss_scale = loss_scale[:pos + 1] + [scale_val] * num_tokens + loss_scale[pos + 1:] + offset += num_tokens + + # Step 4: Compute image_bound using start/end token boundaries + # This is more robust than finding consecutive tokens, especially + # when both image and audio use as placeholder. + input_tensor_ids = torch.tensor(input_ids) + + if images: + im_start_id = tokenizer.convert_tokens_to_ids('') + im_end_id = tokenizer.convert_tokens_to_ids('') + slice_start_id = tokenizer.convert_tokens_to_ids('') + slice_end_id = tokenizer.convert_tokens_to_ids('') + + start_cond = (input_tensor_ids == im_start_id) | (input_tensor_ids == slice_start_id) + end_cond = (input_tensor_ids == im_end_id) | (input_tensor_ids == slice_end_id) + image_start_idx = torch.where(start_cond)[0] + 1 + image_end_idx = torch.where(end_cond)[0] + valid_image_nums = min(len(image_start_idx), len(image_end_idx)) + image_bound = [ + torch.hstack([ + image_start_idx[:valid_image_nums].unsqueeze(-1), + image_end_idx[:valid_image_nums].unsqueeze(-1), + ]) + ] + else: + image_bound = [[]] + + # Step 5: Compute audio_bounds + if inputs.audios: + audio_start_idx = torch.where(input_tensor_ids == audio_start_id)[0] + audio_end_idx = torch.where(input_tensor_ids == audio_end_id)[0] + assert len(audio_start_idx) == len(audio_end_idx) + audio_bounds = [torch.hstack([ + (audio_start_idx + 1).unsqueeze(-1), + audio_end_idx.unsqueeze(-1), + ])] + else: + audio_bounds = [[]] + + encoded = { + 'input_ids': input_ids, + 'labels': labels, + 'loss_scale': loss_scale, + 'image_bound': image_bound, + 'pixel_values': image_inputs['pixel_values'], + 'tgt_sizes': image_inputs['tgt_sizes'], + 'audio_features': audio_features, + 'audio_feature_lens': audio_feature_lens, + 'audio_bounds': audio_bounds, + } + return encoded + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + # Vision embeddings + inputs_embeds, _ = model.get_vllm_embedding(inputs) + # Audio embeddings — scatter audio features into the embedding + inputs_embeds = model.get_omni_embedding( + inputs, + input_embeddings=inputs_embeds, + chunk_length=getattr(self.config, 'audio_chunk_length', 1.0), + ) + return {'inputs_embeds': inputs_embeds} + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = {} + # Vision data + for k in ['pixel_values', 'image_bound', 'tgt_sizes']: + res[k] = self.gather_list(batch, k) + + # Audio data — collect from all samples + all_audio_feats = [] # list of [N_i, 80, frames_i] tensors + audio_feature_lens_list = [] + audio_bounds_list = [] + + for b in batch: + af = b.pop('audio_features', None) + afl = b.pop('audio_feature_lens', None) + ab = b.pop('audio_bounds', None) + + if af is not None and isinstance(af, torch.Tensor): + all_audio_feats.append(af) + if afl is not None: + audio_feature_lens_list.extend(afl) + if ab is not None: + audio_bounds_list.extend(ab) + + # Re-pad audio features across the batch to the same max frame length + if all_audio_feats: + # Unpack per-sample tensors into individual segments, then re-pad + segments = [] + for af in all_audio_feats: + for i in range(af.shape[0]): + segments.append(af[i]) # [80, frames_i] + res['audio_features'] = torch.nn.utils.rnn.pad_sequence( + [s.transpose(0, 1) for s in segments], + batch_first=True, + padding_value=0.0, + ).transpose(1, 2) # [total_segments, 80, max_frames] + else: + res['audio_features'] = [] + + res['audio_feature_lens'] = audio_feature_lens_list if audio_feature_lens_list else [] + res['audio_bounds'] = audio_bounds_list if audio_bounds_list else [] + + res.update(Template._data_collator(self, batch, padding_to=padding_to)) + return res + + +register_template( + Qwen3MixedTemplateMeta( + MLLMTemplateType.minicpmo4_5, + template_cls=MiniCPMO4_5Template, + is_thinking=True, + )) + + +class MiniCPMV4_6Template(Template): + support_padding_free = True + placeholder_tokens = ['<|image_pad|>', '<|video_pad|>'] + + def init_env_args(self): + super().init_env_args() + self.downsample_mode = get_env_args('downsample_mode', str, '16x') + self.max_slice_nums = get_env_args('max_slice_nums', int, 9) + self.video_max_slice_nums = get_env_args('video_max_slice_nums', int, 1) + self.max_num_frames = get_env_args('max_num_frames', int, 128) + self.stack_frames = get_env_args('stack_frames', int, 1) + self.transformers_version = version.parse(transformers.__version__) + self.transformers_5_9 = self.transformers_version >= version.parse('5.9.0') + + def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None: + super()._preprocess_inputs(inputs) + # Inject downsample_mode into mm_processor_kwargs so that vLLM rollout + # receives the correct mode via _encode_truncated -> _add_request. + inputs.mm_processor_kwargs['downsample_mode'] = self.downsample_mode + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + return ['<|image_pad|>\n'] + else: + return ['<|video_pad|>\n'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + split_token = self._tokenize(self.tokenizer.eos_token) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + for media_type in ['image', 'video']: + mm_data = getattr(inputs, f'{media_type}s') + media_token = f'<|{media_type}_pad|>' + media_token_id = self._tokenize(media_token)[0] + max_slice_nums = self.max_slice_nums if media_type == 'image' else self.video_max_slice_nums + if mm_data: + media_inputs = self.processor( + text=self.tokenizer.eos_token.join([media_token] * len(mm_data)), + images=inputs.images or None, + videos=inputs.videos or None, + return_tensors='pt', + add_special_tokens=False, + downsample_mode=self.downsample_mode, + stack_frames=self.stack_frames, + max_num_frames=self.max_num_frames, + max_slice_nums=max_slice_nums, + ) + splited_tokens = self._split_list(media_inputs['input_ids'][0].tolist(), split_token) + idx_list = findall(input_ids, media_token_id) + + def _get_new_tokens(i): + return splited_tokens[i] + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + encoded.update(media_inputs) + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + return encoded + + def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: + res = {} + pixel_values = [b['pixel_values'] for b in batch if b.get('pixel_values') is not None] + if len(pixel_values) > 0: + res['pixel_values'] = torch.concat(pixel_values, dim=-1) + pixel_values_videos = [b['pixel_values_videos'] for b in batch if b.get('pixel_values_videos') is not None] + if len(pixel_values_videos) > 0: + res['pixel_values_videos'] = torch.concat(pixel_values_videos, dim=-1) + + for key in ['target_sizes', 'target_sizes_videos']: + value = self.concat_tensor(batch, key, dim=0) + if value is not None: + res[key] = value + + # Inject downsample_mode so the model forward uses the same mode + # as data preprocessing, keeping image token/feature counts aligned. + res['downsample_mode'] = self.downsample_mode + return res + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + if self.padding_free: + res.update(get_packed_seq_params(res['position_ids'])) + return res + + def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]: + if self.padding_free and self.sequence_parallel_size <= 1 and not self.transformers_5_9: + raise RuntimeError('MiniCPM-V 4.6 packing/padding_free with sequence_parallel_size=1 requires ' + f'transformers>=5.9.0 (current: {self.transformers_version}). ') + return super()._post_encode(model, inputs) + + +register_template( + ChatmlTemplateMeta( + MLLMTemplateType.minicpmv4_6, + template_cls=MiniCPMV4_6Template, + is_thinking=True, + thinking_prefix='\n', + non_thinking_prefix='\n\n\n\n', + )) + +register_template( + ChatmlTemplateMeta( + LLMTemplateType.minicpm5, + is_thinking=True, + thinking_prefix='\n', + non_thinking_prefix='\n\n\n\n', + agent_template='minicpm5', + )) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minimax.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minimax.py new file mode 100644 index 0000000000000000000000000000000000000000..cdac0c190f1b4c89b917d50effc1db0b4af70417 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minimax.py @@ -0,0 +1,301 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from dataclasses import dataclass, field +from typing import Any, Dict, List, Literal, Optional + +from swift.utils import get_env_args, get_logger +from ..base import Template +from ..constant import LLMTemplateType, MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, findall + +logger = get_logger() + + +@dataclass +class MinimaxTemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=list) + prompt: Prompt = field(default_factory=lambda: [ + 'user name=user\n{{QUERY}}\n' + 'ai name=assistant\n' + ]) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['\n']) + suffix: Prompt = field(default_factory=lambda: ['']) + system_prefix: Optional[Prompt] = field( + default_factory=lambda: ['system ai_setting=assistant\n{{SYSTEM}}\n']) + + +register_template(MinimaxTemplateMeta(LLMTemplateType.minimax)) + +register_template( + MinimaxTemplateMeta( + LLMTemplateType.minimax_m1, + prefix=[''], + system_prefix=[ + 'system ai_setting=assistant\n{{SYSTEM}}\n' + ], + )) + + +class MinimaxVLTemplate(Template): + image_placeholder = [''] + skip_prompt = True + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'image' + return self.image_placeholder * inputs.all_image_tokens[index] + + def calc_num_image_tokens(self, image_inputs): + from transformers.image_utils import get_image_size, to_numpy_array + pixel_values = image_inputs['pixel_values'] + image_sizes = image_inputs['image_sizes'] + all_image_tokens = [] + if not image_inputs: + return all_image_tokens + + if self.processor.process_image_mode == 'anyres': + for pixel_value, image_size in zip(pixel_values, image_sizes): + height, width = image_size + num_image_tokens = self.processor.get_num_token(height, width, self.processor.grid_pinpoints, + self.processor.patch_size) + all_image_tokens.append(num_image_tokens) + elif self.processor.process_image_mode == 'resize': + pixel_values = image_inputs['pixel_values'] + all_image_tokens = [] + for pixel_value in pixel_values: + height, width = get_image_size(to_numpy_array(pixel_value)) + all_image_tokens.append(int(height * width / self.processor.patch_size**2)) + else: + if self.processor.patch_size is not None: + pixel_values = image_inputs['pixel_values'] + all_image_tokens = [] + for pixel_value in pixel_values: + height, width = get_image_size(to_numpy_array(pixel_value)) + new_width, new_height = self.processor.get_hw_multiple_of( + (width, height), self.processor.patch_size, self.processor.max_size) + num_image_tokens = ((new_height // self.processor.patch_size) * + (new_width // self.processor.patch_size)) # + 1 + all_image_tokens.append(num_image_tokens) + else: + logger.warning_once( + 'Expanding inputs for image tokens in MiniMaxVL01 should be done in processing. ' + "Please add `patch_size` and `vision_feature_select_strategy` to the model's " + 'processing config or set directly ' + 'with `processor.patch_size = {{patch_size}}` and processor.vision_feature_select_strategy = ' + '{{vision_feature_select_strategy}}`. ' + 'Using processors without these attributes in the config is deprecated ' + 'and will throw an error in v4.47.') + raise ValueError( + "You need to provide `patch_size` and `vision_feature_select_strategy` in the model's processing " + 'config to expand inputs for image tokens.') + return all_image_tokens + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + output_kwargs = self.processor._merge_kwargs( + self.processor.MiniMaxVL01ProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + ) + if inputs.images: + image_inputs = self.processor.image_processor( + inputs.images, **output_kwargs['images_kwargs'], return_tensors='pt') + inputs.all_image_tokens = self.calc_num_image_tokens(image_inputs) + else: + image_inputs = {} + encoded = super()._encode(inputs) + for key in image_inputs: + encoded[key] = image_inputs[key] + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + pixel_values = self.gather_list(batch, 'pixel_values') + image_sizes = self.gather_list(batch, 'image_sizes') + res = super()._data_collator(batch, padding_to=padding_to) + if pixel_values: + res['pixel_values'] = pixel_values + if image_sizes: + res['image_sizes'] = image_sizes + return res + + +register_template(MinimaxTemplateMeta(LLMTemplateType.minimax_vl, template_cls=MinimaxVLTemplate)) + + +@dataclass +class MinimaxM2TemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=lambda: [']~!b[]~b]system\n{{SYSTEM}}[e~[\n']) + prompt: Prompt = field(default_factory=lambda: [']~b]user\n{{QUERY}}[e~[\n]~b]ai\n']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['[e~[\n']) + suffix: Prompt = field(default_factory=lambda: ['[e~[\n']) + agent_template: Optional[str] = 'minimax_m2' + is_thinking: bool = True + thinking_prefix: str = '\n' + + +register_template( + MinimaxM2TemplateMeta( + LLMTemplateType.minimax_m2, + default_system='You are MiniMax-M2, a helpful AI assistant built by MiniMax. Knowledge cutoff: 2025-06.', + )) + +register_template( + MinimaxM2TemplateMeta( + LLMTemplateType.minimax_m2_1, + default_system='You are a helpful assistant. Your name is MiniMax-M2.1 and is built by MiniMax.', + )) + +register_template( + MinimaxM2TemplateMeta( + LLMTemplateType.minimax_m2_5, + default_system='You are a helpful assistant. Your name is MiniMax-M2.5 and is built by MiniMax.', + )) + +register_template( + MinimaxM2TemplateMeta( + LLMTemplateType.minimax_m2_7, + default_system='You are a helpful assistant. Your name is MiniMax-M2.7 and is built by MiniMax.', + )) + +_MINIMAX_M3_IDENTITY = ('Your model version is MiniMax-M3, developed by MiniMax. Knowledge cutoff: January 2026. ' + 'Founded in early 2022, MiniMax is a global AI foundation model company committed to ' + 'advancing the frontiers of AI towards AGI.') + +_MINIMAX_M3_THINKING_BASE = ( + 'You have a thinking capability that allows you to reason step by step before responding. ' + 'When thinking is enabled, wrap your reasoning in tags before your ' + 'response. When thinking is disabled, begin your response directly after the ' + 'prefix. When thinking is adaptive, decide on your own whether to think for the current turn.') + +_MINIMAX_M3_THINKING_MODE_TEXT = { + 'enabled': ('Current thinking mode: enabled. You MUST think step by step before every response, ' + 'including after receiving function/tool results.'), + 'disabled': + 'Current thinking mode: disabled. Do not output any thinking process.', + 'adaptive': ('Current thinking mode: adaptive. You are encouraged to think for complex ' + 'decision-making, multi-step reasoning, or when analyzing function/tool results.'), +} + +_MINIMAX_M3_DEFAULT_DEVELOPER = 'You are a helpful assistant.' + + +def _build_m3_system_block(thinking_mode: str = 'adaptive') -> str: + mode_text = _MINIMAX_M3_THINKING_MODE_TEXT.get(thinking_mode, _MINIMAX_M3_THINKING_MODE_TEXT['adaptive']) + return (f'{_MINIMAX_M3_IDENTITY}' + f'\n\n\n{_MINIMAX_M3_THINKING_BASE}\n{mode_text}\n') + + +class MinimaxM3VLTemplate(Template): + image_token = ']<]image[>[' + video_token = ']<]video[>[' + placeholder_tokens = [']<]image[>[', ']<]video[>['] + + def init_env_args(self): + super().init_env_args() + # thinking_mode: "enabled" / "disabled" / "adaptive" + self.thinking_mode = get_env_args('thinking_mode', str, 'disabled') + self.chat_template_kwargs['thinking_mode'] = self.thinking_mode + # Map thinking_mode to enable_thinking for the broader framework + if self.thinking_mode == 'disabled': + self.enable_thinking = False + else: + self.enable_thinking = True + + def _get_thinking_mode(self, inputs=None) -> str: + thinking_mode = None if inputs is None else inputs.chat_template_kwargs.get('thinking_mode') + if thinking_mode is None: + thinking_mode = self.chat_template_kwargs.get('thinking_mode', 'adaptive') + return thinking_mode + + def _get_enable_thinking(self, inputs=None): + thinking_mode = self._get_thinking_mode(inputs) + return thinking_mode != 'disabled' + + def _get_response_prefix(self, inputs=None): + # Check explicit override first + response_prefix = None if inputs is None else inputs.chat_template_kwargs.get('response_prefix') + if response_prefix is not None: + return response_prefix + if self.response_prefix is not None: + return self.response_prefix + thinking_mode = self._get_thinking_mode(inputs) + if thinking_mode == 'enabled': + return self.template_meta.thinking_prefix # '' + elif thinking_mode == 'disabled': + return self.template_meta.non_thinking_prefix # '' + else: # adaptive + return '' # No prefix, let model decide + + def _get_system(self, inputs: StdTemplateInputs) -> str: + system = super()._get_system(inputs) + thinking_mode = self._get_thinking_mode(inputs) + system_block = _build_m3_system_block(thinking_mode) + return f'{system_block}[e~[\n]~b]developer\n{system or ""}' + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + return [self.image_token] + elif media_type == 'video': + return [self.video_token] + else: + raise ValueError(f'Unsupported media type for MiniMax-M3 VL: {media_type}') + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + if not inputs.images and not inputs.videos: + return encoded + media_text_parts = ([self.image_token] * len(inputs.images) + [self.video_token] * len(inputs.videos)) + media_inputs = self.processor( + text=self.tokenizer.eos_token.join(media_text_parts), + images=inputs.images or None, + videos=inputs.videos or None, + return_tensors='pt', + ) + split_token = self._tokenize(self.tokenizer.eos_token) + splited_tokens = self._split_list(media_inputs['input_ids'][0].tolist(), split_token) + media_inputs.pop('input_ids', None) + media_inputs.pop('attention_mask', None) + + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + + idx_list = [] + for key in ['image', 'video']: + token_id = getattr(self.config, f'{key}_token_id', None) + if token_id is None: + continue + idx_list += findall(input_ids, token_id) + sorted_order = sorted(range(len(idx_list)), key=lambda i: idx_list[i]) + idx_list = [idx_list[i] for i in sorted_order] + splited_tokens = [splited_tokens[i] for i in sorted_order] + + def _get_new_tokens(i): + return splited_tokens[i] + + if idx_list: + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + encoded.update(media_inputs) + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + return encoded + + +@dataclass +class MinimaxM3VLTemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=lambda: [']~!b[]~b]system\n{{SYSTEM}}[e~[\n']) + prompt: Prompt = field(default_factory=lambda: [']~b]user\n{{QUERY}}[e~[\n]~b]ai\n']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['[e~[\n']) + suffix: Prompt = field(default_factory=lambda: ['[e~[\n']) + default_system: Optional[str] = _MINIMAX_M3_DEFAULT_DEVELOPER + agent_template: Optional[str] = 'minimax_m3' + is_thinking: bool = True + thinking_prefix: str = '' + non_thinking_prefix: str = '' + history_thinking_prefix: str = '' + + +register_template(MinimaxM3VLTemplateMeta(MLLMTemplateType.minimax_m3_vl, template_cls=MinimaxM3VLTemplate)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minimind.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minimind.py new file mode 100644 index 0000000000000000000000000000000000000000..60f9a4737839332c7d51248db9a37aebdc3577a5 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/minimind.py @@ -0,0 +1,22 @@ +from dataclasses import dataclass, field +from typing import List, Optional + +from ..constant import LLMTemplateType +from ..register import TemplateMeta, register_template +from ..utils import Prompt, Word + +DEFAULT_SYSTEM = 'You are a helpful assistant' + + +@dataclass +class MiniMindTemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=list) + prompt: Prompt = field(default_factory=lambda: ['<|im_start|>user\n{{QUERY}}<|im_end|>\n<|im_start|>assistant\n']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|im_end|>\n']) + suffix: Prompt = field(default_factory=lambda: ['<|im_end|>\n']) + system_prefix: Optional[Prompt] = field(default_factory=lambda: ['<|im_start|>system\n{{SYSTEM}}<|im_end|>\n']) + default_system: Optional[str] = DEFAULT_SYSTEM + stop_words: List[Word] = field(default_factory=lambda: ['<|endoftext|>']) + + +register_template(MiniMindTemplateMeta(LLMTemplateType.minimind)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mistral.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mistral.py new file mode 100644 index 0000000000000000000000000000000000000000..f737f8efc80f3df91f2e3182b636236be7a24c98 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mistral.py @@ -0,0 +1,203 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import os +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from typing import Any, Dict, List, Literal, Optional + +from ..base import Template +from ..constant import LLMTemplateType, MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, findall + +today = datetime.now().strftime('%Y-%m-%d') + +mistral_2501_system = ( + 'You are Mistral Small 3, a Large Language Model (LLM) created by Mistral AI, a French startup ' + 'headquartered in Paris.\n' + f'Your knowledge base was last updated on 2023-10-01. The current date is {today}.\n\n' + "When you're not sure about some information, you say that you don't have the information and don't " + 'make up anything.\n' + "If the user's question is not clear, ambiguous, or does not provide enough context for you to accurately answer " + 'the question, you do not try to answer it right away and you rather ask the user to clarify their request (e.g. ' + '"What are some good restaurants around me?" => "Where are you?" or "When is the next flight to Tokyo" => "' + 'Where do you travel from?")') + + +@dataclass +class Mistral3TemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=lambda: ['']) + prompt: Prompt = field(default_factory=lambda: ['[INST]{{QUERY}}[/INST]']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['']) + suffix: Prompt = field(default_factory=lambda: ['']) + system_prefix: Optional[Prompt] = field(default_factory=lambda: ['[SYSTEM_PROMPT]{{SYSTEM}}[/SYSTEM_PROMPT]']) + + +register_template(Mistral3TemplateMeta(LLMTemplateType.mistral_2501, default_system=mistral_2501_system)) + + +class Mistral2503Template(Template): + placeholder_tokens = ['[IMG]'] + image_token = 10 + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'image' + return ['[IMG]'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + processor = self.processor + images = inputs.images + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + idx_list = findall(input_ids, self.image_token) + patch_size = processor.patch_size * processor.spatial_merge_size + if idx_list: + image_inputs = processor.image_processor(images, patch_size=patch_size, return_tensors='pt') + encoded['pixel_values'] = image_inputs['pixel_values'].to(self.model_info.torch_dtype) + encoded['image_sizes'] = image_sizes = image_inputs['image_sizes'] + + def _get_new_tokens(i): + height, width = image_sizes[i] + num_height_tokens = height // patch_size + num_width_tokens = width // patch_size + replace_tokens = [[processor.image_token] * num_width_tokens + [processor.image_break_token] + ] * num_height_tokens + # Flatten list + replace_tokens = [item for sublist in replace_tokens for item in sublist] + replace_tokens[-1] = processor.image_end_token + replace_str = ''.join(replace_tokens) + return processor.encode(replace_str, add_special_tokens=False) + + encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens( + input_ids, labels, loss_scale, idx_list, _get_new_tokens) + + return encoded + + +register_template( + Mistral3TemplateMeta( + MLLMTemplateType.mistral_2503, default_system=mistral_2501_system, template_cls=Mistral2503Template)) + +devstral_small_2505_system = ( # from https://huggingface.co/mistralai/Devstral-Small-2505/blob/main/SYSTEM_PROMPT.txt + 'You are Devstral, a helpful agentic model trained by Mistral AI and using the OpenHands scaffold. ' + 'You can interact with a computer to solve tasks.\n\n\nYour primary role is to assist users by ' + 'executing commands, modifying code, and solving technical problems effectively. You should be ' + 'thorough, methodical, and prioritize quality over speed.\n* If the user asks a question, like ' + '"why is X happening", don\'t try to fix the problem. Just give an answer to the question.' + '\n\n\n\n* Each action you take is somewhat expensive. Wherever possible, ' + 'combine multiple actions into a single action, e.g. combine multiple bash commands into one, using ' + 'sed and grep to edit/view multiple files at once.\n* When exploring the codebase, use efficient tools ' + 'like find, grep, and git commands with appropriate filters to minimize unnecessary operations.' + '\n\n\n\n* When a user provides a file path, do NOT assume it\'s ' + 'relative to the current working directory. First explore the file system to locate the file before ' + 'working on it.\n* If asked to edit a file, edit the file directly, rather than creating a new file with ' + 'a different filename.\n* For global search-and-replace operations, consider using `sed` instead of ' + 'opening file editors multiple times.\n\n\n\n* Write clean, ' + 'efficient code with minimal comments. Avoid redundancy in comments: Do not repeat information that can ' + 'be easily inferred from the code itself.\n* When implementing solutions, focus on making the minimal ' + 'changes needed to solve the problem.\n* Before implementing any changes, first thoroughly understand ' + 'the codebase through exploration.\n* If you are adding a lot of code to a function or file, consider ' + 'splitting the function or file into smaller pieces when appropriate.\n\n\n' + '\n* When configuring git credentials, use "openhands" as the user.name and ' + '"openhands@all-hands.dev" as the user.email by default, unless explicitly instructed otherwise.' + '\n* Exercise caution with git operations. Do NOT make potentially dangerous changes (e.g., pushing ' + 'to main, deleting repositories) unless explicitly asked to do so.\n* When committing changes, use `git' + ' status` to see all modified files, and stage all files necessary for the commit. Use `git commit -a` ' + 'whenever possible.\n* Do NOT commit files that typically shouldn\'t go into version control (e.g., ' + 'node_modules/, .env files, build directories, cache files, large binaries) unless explicitly ' + 'instructed by the user.\n* If unsure about committing certain files, check for the presence of .' + 'gitignore files or ask the user for clarification.\n\n\n\n* When ' + 'creating pull requests, create only ONE per session/issue unless explicitly instructed otherwise.\n* ' + 'When working with an existing PR, update it with new commits rather than creating additional PRs for ' + 'the same issue.\n* When updating a PR, preserve the original PR title and purpose, updating description ' + 'only when necessary.\n\n\n\n1. EXPLORATION: Thoroughly ' + 'explore relevant files and understand the context before proposing solutions\n2. ANALYSIS: Consider ' + 'multiple approaches and select the most promising one\n3. TESTING:\n * For bug fixes: Create tests to ' + 'verify issues before implementing fixes\n * For new features: Consider test-driven development when ' + 'appropriate\n * If the repository lacks testing infrastructure and implementing tests would require ' + 'extensive setup, consult with the user before investing time in building testing infrastructure\n * ' + 'If the environment is not set up to run tests, consult with the user first before investing time to ' + 'install all dependencies\n4. IMPLEMENTATION: Make focused, minimal changes to address the problem\n5. ' + 'VERIFICATION: If the environment is set up to run tests, test your implementation thoroughly, including ' + 'edge cases. If the environment is not set up to run tests, consult with the user first before investing ' + 'time to run tests.\n\n\n\n* Only use GITHUB_TOKEN and other ' + 'credentials in ways the user has explicitly requested and would expect.\n* Use APIs to work with GitHub ' + 'or other platforms, unless the user asks otherwise or your task requires browsing.\n\n\n\n* When user asks you to run an application, don\'t stop if the ' + 'application is not installed. Instead, please install the application and run the command again.\n* If ' + 'you encounter missing dependencies:\n 1. First, look around in the repository for existing dependency ' + 'files (requirements.txt, pyproject.toml, package.json, Gemfile, etc.)\n 2. If dependency files exist, ' + 'use them to install all dependencies at once (e.g., `pip install -r requirements.txt`, `npm install`, ' + 'etc.)\n 3. Only install individual packages directly if no dependency files are found or if only ' + 'specific packages are needed\n* Similarly, if you encounter missing dependencies for essential tools ' + 'requested by the user, install them when possible.\n\n\n\n* If ' + 'you\'ve made repeated attempts to solve a problem but tests still fail or the user reports it\'s still ' + 'broken:\n 1. Step back and reflect on 5-7 different possible sources of the problem\n 2. Assess the ' + 'likelihood of each possible cause\n 3. Methodically address the most likely causes, starting with the ' + 'highest probability\n 4. Document your reasoning process\n* When you run into any major issue while ' + 'executing a plan from the user, please don\'t try to directly work around it. Instead, propose a new ' + 'plan and confirm with the user before proceeding.\n') + +register_template(Mistral3TemplateMeta('devstral', default_system=devstral_small_2505_system)) + + +class Mistral2506Template(Mistral2503Template): + + def _get_mistral_system(self): + from swift.model import get_model_name + model_dir = self.model_info.model_dir + model_name = get_model_name(model_dir) + file_path = os.path.join(model_dir, 'SYSTEM_PROMPT.txt') + with open(file_path, 'r', encoding='utf-8') as file: + system_prompt = file.read() + today = datetime.today().strftime('%Y-%m-%d') + yesterday = (datetime.today() - timedelta(days=1)).strftime('%Y-%m-%d') + return system_prompt.format(name=model_name, today=today, yesterday=yesterday) + + def _get_system(self, inputs: StdTemplateInputs) -> Optional[str]: + system = super()._get_system(inputs) + if system is None: + system = self._get_mistral_system() + return system + + +register_template( + Mistral3TemplateMeta(MLLMTemplateType.mistral_2506, default_system=None, template_cls=Mistral2506Template)) + + +class Mistral2512Template(Mistral2506Template): + + def _get_mistral_system(self): + model_dir = self.model_info.model_dir + file_path = os.path.join(model_dir, 'SYSTEM_PROMPT.txt') + with open(file_path, 'r', encoding='utf-8') as file: + system_prompt = file.read() + today = datetime.today().strftime('%Y-%m-%d') + yesterday = (datetime.today() - timedelta(days=1)).strftime('%Y-%m-%d') + return system_prompt.format(today=today, yesterday=yesterday) + + +register_template( + Mistral3TemplateMeta( + MLLMTemplateType.mistral_2512, default_system=None, template_cls=Mistral2512Template, agent_template='mistral')) + + +class Mistral2512ThinkingTemplate(Mistral2506Template): + + def _get_mistral_system(self): + model_dir = self.model_info.model_dir + file_path = os.path.join(model_dir, 'SYSTEM_PROMPT.txt') + with open(file_path, 'r', encoding='utf-8') as file: + system_prompt = file.read() + return system_prompt + + +register_template( + Mistral3TemplateMeta( + MLLMTemplateType.mistral_2512_thinking, + default_system=None, + template_cls=Mistral2512ThinkingTemplate, + agent_template='mistral')) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/molmo.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/molmo.py new file mode 100644 index 0000000000000000000000000000000000000000..d0140ef69a4eb5c49c91da036d38ac9d207ede0e --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/molmo.py @@ -0,0 +1,177 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import numpy as np +import torch +from PIL import Image +from typing import Any, Dict, List, Literal, Optional, Tuple + +from ..base import Template +from ..constant import MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, findall +from .utils import ChatmlTemplateMeta + + +class MolmoTemplate(Template): + placeholder_tokens = [''] + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + return [] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + # image + images_inputs = self.processor.process(images=inputs.images or None, text='') + images_input_ids = images_inputs.pop('input_ids').tolist() + user_token = self._tokenize(' User') + assert len(user_token) == 1 + idx = findall(images_input_ids, user_token[0]) + assert len(idx) == 1 + labels = encoded['labels'] + encoded['input_ids'] = images_input_ids[:idx[0]] + encoded['input_ids'] + if labels: + encoded['labels'] = [-100] * idx[0] + labels + if 'images' in images_inputs: + images_inputs['images'] = images_inputs['images'].to(self.model_info.torch_dtype) + encoded.update(images_inputs) + return encoded + + def generate(self, model, **kwargs): + kwargs.pop('attention_mask', None) + generation_config = kwargs.pop('generation_config') + batch = { + k: kwargs.pop(k, None) + for k in ['input_ids', 'attention_mask', 'images', 'image_input_idx', 'image_masks'] + } + return model.generate_from_batch(batch, generation_config, **kwargs) + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + # prepare batchfy inputs + keys = ['images', 'image_input_idx', 'image_masks'] + images_res = self.fetch_inputs(batch, keys) + for key in keys: + val = images_res.get(key) + if val: + images_res[key] = torch.stack(val) + res.update(images_res) + return res + + +register_template( + TemplateMeta( + MLLMTemplateType.molmo, + prefix=[], + prompt=[' User: {{QUERY}} Assistant:'], + chat_sep=None, + suffix=['<|endoftext|>'], + template_cls=MolmoTemplate, + )) + + +class Molmo2Template(Template): + """Molmo2 template for image and video understanding. + + Uses ChatML format with BOS auto-insertion. + Media placeholders (<|image|>, <|video|>) are expanded via _extend_tokens. + Video loading/sampling is delegated entirely to processor.video_processor. + """ + + use_model = True + + placeholder_tokens = [ + '<|image|>', + '<|video|>', + '', + '', + '', + ] + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + return ['<|image|>'] + elif media_type == 'video': + return ['<|video|>'] + else: + raise ValueError(f'Unsupported media_type: {media_type}') + + def _prepare_mm_inputs(self, inputs: StdTemplateInputs) -> Tuple[Dict[str, Any], List[List[int]], List[List[int]]]: + media_inputs: Dict[str, Any] = {} + image_expansions: List[List[int]] = [] + video_expansions: List[List[int]] = [] + tokenizer = self.tokenizer + + if inputs.images: + image_inputs = self.processor.image_processor(images=inputs.images, return_tensors='pt') + for image_grid in image_inputs['image_grids']: + image_tokens = self.processor.get_image_tokens(image_grid.cpu().numpy()) + image_expansions.append(tokenizer.encode(''.join(image_tokens), add_special_tokens=False)) + media_inputs.update(image_inputs) + + if inputs.videos: + if len(inputs.videos) != 1: + raise ValueError('Molmo2 currently only supports single-video samples.') + video_inputs = self.processor.video_processor( + videos=inputs.videos, + return_tensors='pt', + return_metadata=True, + ) + video_metadata = video_inputs.pop('video_metadata') + for video_grid, metadata in zip(video_inputs['video_grids'], video_metadata): + video_string = self.processor.get_video_string( + video_grid.cpu().numpy(), + np.asarray(metadata.timestamps, dtype=np.float32), + ) + video_expansions.append(tokenizer.encode(video_string, add_special_tokens=False)) + media_inputs.update(video_inputs) + + return media_inputs, image_expansions, video_expansions + + def _build_token_type_ids(self, input_ids: List[int]) -> List[int]: + image_token_ids = {int(token_id) for token_id in getattr(self.processor, 'image_token_ids', [])} + return [1 if token_id in image_token_ids else 0 for token_id in input_ids] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + media_inputs, image_expansions, video_expansions = self._prepare_mm_inputs(inputs) + + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale') + + # Expand image placeholders + image_placeholder = self.tokenizer.convert_tokens_to_ids('<|image|>') + idx_list = findall(input_ids, image_placeholder) + if idx_list: + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + lambda i: image_expansions[i]) + + # Expand video placeholders + video_placeholder = self.tokenizer.convert_tokens_to_ids('<|video|>') + idx_list = findall(input_ids, video_placeholder) + if idx_list: + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + lambda i: video_expansions[i]) + + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + encoded['token_type_ids'] = self._build_token_type_ids(input_ids) + encoded.update(media_inputs) + return encoded + + def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: + res = super()._data_collator_mm_data(batch) + for key in ['image_grids', 'video_grids', 'image_token_pooling', 'video_token_pooling', 'image_num_crops']: + value = self.concat_tensor(batch, key, 0) + if value is not None: + res[key] = value + return res + + +register_template(ChatmlTemplateMeta( + MLLMTemplateType.molmo2, + template_cls=Molmo2Template, +)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/moonshot.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/moonshot.py new file mode 100644 index 0000000000000000000000000000000000000000..6cba580c9ca036e16b42b331bd68eb66180d3c0d --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/moonshot.py @@ -0,0 +1,439 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +import math +import torch +from dataclasses import dataclass, field +from PIL import Image +from torch import nn as nn +from typing import Any, Dict, List, Literal, Optional + +from swift.utils import is_deepspeed_enabled, to_device +from ..base import Template +from ..constant import LLMTemplateType, MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, findall + + +@dataclass +class MoonlightTemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=list) + prompt: Prompt = field(default_factory=lambda: + ['<|im_user|>user<|im_middle|>{{QUERY}}<|im_end|><|im_assistant|>assistant<|im_middle|>']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|im_end|>']) + suffix: Prompt = field(default_factory=lambda: ['<|im_end|>']) + system_prefix: Optional[Prompt] = field( + default_factory=lambda: ['<|im_system|>system<|im_middle|>{{SYSTEM}}<|im_end|>']) + default_system: Optional[str] = 'You are a helpful assistant' + + +register_template(MoonlightTemplateMeta(LLMTemplateType.moonlight)) + +register_template( + MoonlightTemplateMeta( + LLMTemplateType.kimi_k2, default_system='You are Kimi, an AI assistant created by Moonshot AI.')) + + +class KimiVLTemplate(Template): + placeholder_tokens = ['<|media_pad|>'] + support_padding_free = True + skip_prompt = False + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + return ['<|media_start|>image<|media_content|><|media_pad|><|media_end|>'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + media_token = self._tokenize('<|media_pad|>')[0] + idx_list = findall(input_ids, media_token) + if inputs.images: + image_processor = self.processor.image_processor + image_inputs = image_processor(inputs.images, return_tensors='pt') + image_grid_hws = image_inputs['image_grid_hws'] + merge_length = image_processor.merge_kernel_size[0] * image_processor.merge_kernel_size[1] + + def _get_new_tokens(i): + token_len = (image_grid_hws[i].prod() // merge_length) + return [media_token] * token_len + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + + encoded['loss_scale'] = loss_scale + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded.update(image_inputs) + return encoded + + def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: + res = super()._data_collator_mm_data(batch) + image_grid_hws = self.concat_tensor(batch, 'image_grid_hws', 0) + if image_grid_hws is not None: + res['image_grid_hws'] = image_grid_hws + return res + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + input_ids = inputs['input_ids'] + pixel_values = inputs.get('pixel_values') + inputs_embeds = model.get_input_embeddings()(input_ids) + + if pixel_values is not None and pixel_values.size(0) > 0: + pixel_values = pixel_values.to(model.vision_tower.dtype) + image_features: torch.Tensor = model._extract_image_features(pixel_values, inputs['image_grid_hws']) + inputs_embeds = inputs_embeds.to(image_features[0].dtype).clone() + inputs_embeds = model._merge_with_image_features(inputs_embeds, input_ids, image_features) + elif is_deepspeed_enabled(): + image_processor = self.processor.image_processor + dummy_image = Image.new('RGB', (32, 32), (0, 0, 0)) + image_inputs = image_processor([dummy_image], return_tensors='pt') + pixel_values = image_inputs['pixel_values'].to(model.vision_tower.dtype) + image_features: torch.Tensor = model._extract_image_features(pixel_values, image_inputs['image_grid_hws']) + inputs_embeds = inputs_embeds + image_features.mean() * 0. + return {'inputs_embeds': inputs_embeds} + + +register_template(MoonlightTemplateMeta(MLLMTemplateType.kimi_vl, template_cls=KimiVLTemplate)) + + +class KimiK25Template(Template): + placeholder_tokens = ['<|media_pad|>', '<|kimi_k25_video_placeholder|>'] + jinja_enable_thinking_key = 'thinking' + support_padding_free = True + skip_prompt = False + + def _get_system(self, inputs: StdTemplateInputs) -> Optional[str]: + system = super()._get_system(inputs) + if system is not None and '<|im_middle|>' not in system: # compat agent + system = f'system<|im_middle|>{system}' + return system + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + return ['<|media_begin|>image<|media_content|><|media_pad|><|media_end|>\n'] + raise ValueError(f'KimiK25Template does not currently support {media_type}. ' + 'Please open an issue to request support.') + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + media_token = self._tokenize('<|media_pad|>')[0] + idx_list = findall(input_ids, media_token) + if inputs.images: + image_processor = self.processor.image_processor + image_inputs = image_processor([{ + 'type': 'image', + 'image': image + } for image in inputs.images], + return_tensors='pt') + grid_thws = image_inputs['grid_thws'] + merge_length = math.prod(self.config.vision_config.merge_kernel_size) + + def _get_new_tokens(i): + token_len = (grid_thws[i].prod() // merge_length) + return [media_token] * token_len + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + + encoded['loss_scale'] = loss_scale + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded.update(image_inputs) + return encoded + + def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: + res = super()._data_collator_mm_data(batch) + grid_thws = self.concat_tensor(batch, 'grid_thws', 0) + if grid_thws is not None: + res['grid_thws'] = grid_thws + return res + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + input_ids = inputs['input_ids'] + pixel_values = inputs.get('pixel_values') + inputs_embeds = model.get_input_embeddings()(input_ids) + + if pixel_values is not None and pixel_values.size(0) > 0: + pixel_values = pixel_values.to(model.vision_tower.dtype) + image_features: torch.Tensor = model._extract_image_features(pixel_values, inputs['grid_thws']) + if model.mm_projector: + image_features = model.mm_projector(image_features) + image_features = torch.cat(image_features, dim=0) + inputs_embeds = inputs_embeds.to(image_features.dtype) + image_mask = (input_ids == self.config.media_placeholder_token_id).unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_features) + elif is_deepspeed_enabled(): + image_processor = self.processor.image_processor + dummy_image = Image.new('RGB', (32, 32), (0, 0, 0)) + image_inputs = image_processor([{'type': 'image', 'image': dummy_image}], return_tensors='pt') + image_inputs = to_device(image_inputs, inputs_embeds.device) + pixel_values = image_inputs['pixel_values'].to(model.vision_tower.dtype) + image_features: torch.Tensor = model._extract_image_features(pixel_values, image_inputs['grid_thws']) + if model.mm_projector: + image_features = model.mm_projector(image_features) + image_features = torch.cat(image_features, dim=0) + inputs_embeds = inputs_embeds + image_features.mean() * 0. + return {'inputs_embeds': inputs_embeds} + + +register_template( + MoonlightTemplateMeta( + MLLMTemplateType.kimi_k25, + template_cls=KimiK25Template, + system_prefix=['<|im_system|>{{SYSTEM}}<|im_end|>'], + default_system=None, + is_thinking=True, + thinking_prefix='', + non_thinking_prefix='', + history_thinking_prefix='', + agent_template='kimi_k25', + )) + + +class KimiK3Template(Template): + # Kimi-K3 renders chats in XTML: structural markers (<|open|>/<|close|>/<|sep|>/<|end_of_msg|>) + # are special tokens while tag names (message/think/response) are plain text. + # See `encoding_k3.py` in the model repo. + placeholder_tokens = ['<|media_pad|>'] + support_padding_free = True + skip_prompt = False + jinja_enable_thinking_key = 'thinking' + + think_open = '<|open|>think<|sep|>' + think_close = '<|close|>think<|sep|>' + response_open = '<|open|>response<|sep|>' + response_close = '<|close|>response<|sep|>' + valid_thinking_efforts = {'low', 'high', 'max'} + + def __init__(self, *args, **kwargs): + # Kimi K3 always has thinking enabled (the base class defaults hybrid-thinking + # templates with a non_thinking_prefix to False). + if kwargs.get('enable_thinking') is None: + kwargs['enable_thinking'] = True + super().__init__(*args, **kwargs) + + def _thinking_to_xtml(self, text: str, complete: bool = True) -> str: + """Convert the inline `...` convention into the K3 think/response channels. + + The channel is structural: every assistant message carries the open/close + tags even when there is no reasoning content (aligned with `encoding_k3.py`). + A complete message also owns the response-channel close, so that a tool-calls + section can follow it (chat_sep/suffix only close the enclosing message). + """ + if text.startswith(self.think_open): # already converted + return text + if text.startswith(''): + body = text[len(''):] + think, sep, response = body.partition('') + if not sep: # bare thinking prefix (generation prompt) + return self.think_open + think + res = self.think_open + think + self.think_close + self.response_open + response + else: + res = self.think_open + self.think_close + self.response_open + text + if complete: + res += self.response_close + return res + + def _xtml_to_thinking(self, text: str) -> str: + """Inverse of `_thinking_to_xtml`: map the K3 think/response channels back to + the inline `...` convention used across swift, so that decoded + responses round-trip through `_preprocess_inputs` in multi-turn conversations. + A tool-calls section (`<|open|>tools<|sep|>...`) is left untouched for + `get_toolcall` to parse. + """ + text = text.replace(self.think_close + self.response_open, '') + text = text.replace(self.think_open, '') + # Unpaired markers (e.g. truncated or non-standard generations). + text = text.replace(self.think_close, '') + text = text.replace(self.response_open, '') + return text.replace(self.response_close, '') + + def decode_generate_ids(self, generate_ids, **kwargs): + response = super().decode_generate_ids(generate_ids, **kwargs) + if isinstance(response, str): + response = self._xtml_to_thinking(response) + return response + + def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None: + super()._preprocess_inputs(inputs) + # Convert every assistant message into the XTML think/response channel form + # ahead of the base thinking machinery, so that the rendered response always + # equals the message content (required by loss_scale). + for message in inputs.messages: + if message['role'] == 'assistant' and isinstance(message['content'], str): + message['content'] = self._thinking_to_xtml(message['content']) + + def _add_non_thinking_prefix(self, inputs, thinking_prefix='') -> None: + return super()._add_non_thinking_prefix(inputs, thinking_prefix=self.think_open) + + def _remove_thinking_content(self, content: str, thinking_suffix='') -> str: + content = content.split(self.think_close)[-1].strip() + if content.startswith(self.response_open): + content = content[len(self.response_open):] + return self.template_meta.history_thinking_prefix + content + + def _get_preserve_thinking(self, inputs=None): + preserve_thinking = None if inputs is None else inputs.chat_template_kwargs.get('preserve_thinking') + if preserve_thinking is None: + preserve_thinking = self.preserve_thinking + if preserve_thinking is None: + # K3 was trained in preserved-thinking-history mode: keep historical + # reasoning by default (pass `--preserve_thinking false` to drop it). + preserve_thinking = True + return preserve_thinking + + def _get_system(self, inputs: StdTemplateInputs) -> Optional[str]: + system = super()._get_system(inputs) + if system is not None and '<|sep|>' not in system: # compat agent + system = f'role="system"<|sep|>{system}' + return system + + def _get_response_prefix(self, inputs=None): + response_prefix = super()._get_response_prefix(inputs) + if not response_prefix: + return response_prefix + # A user-supplied response_prefix may still use the inline convention; + # a generation prefix must keep the response channel open. + return self._thinking_to_xtml(response_prefix, complete=False) + + def _get_thinking_effort(self, inputs=None): + # `reasoning_effort` is the K3 API field name; `thinking_effort` is used by the + # chat encoder (encoding_k3.py). Accept both via chat_template_kwargs. + kwargs = {} if inputs is None else inputs.chat_template_kwargs + thinking_effort = kwargs.get('thinking_effort') or kwargs.get('reasoning_effort') + if thinking_effort is not None: + assert thinking_effort in self.valid_thinking_efforts, ( + f'Unsupported thinking_effort={thinking_effort!r}; ' + f'supported values are {sorted(self.valid_thinking_efforts)}.') + return thinking_effort + + def _swift_encode(self, inputs: StdTemplateInputs): + res_context_list, loss_scale_list, answer_len = super()._swift_encode(inputs) + thinking_effort = self._get_thinking_effort(inputs) + if thinking_effort is not None: + # Aligned with `_internal_system_message` in encoding_k3.py: the + # thinking-effort system message is rendered before all messages, + # but after the tool-declare system message when tools are present. + context = ('<|open|>message role="system" type="thinking-effort"<|sep|>' + '`thinking_effort` guides on how much to think in your ' + 'thinking channel (not including the response channel), ' + 'supported values include `low`, `medium`, `high`, and `max`.\n' + f'Now the system is invoked with `thinking_effort={thinking_effort}`.' + '<|close|>message<|sep|><|end_of_msg|>') + first = res_context_list[0] if res_context_list else None + if inputs.tools and isinstance(first, str) and 'type="tool-declare"' in first: + end = first.index('<|end_of_msg|>') + len('<|end_of_msg|>') + res_context_list[0] = first[:end] + context + first[end:] + else: + res_context_list.insert(0, context) + loss_scale_list.insert(0, 0.) + return res_context_list, loss_scale_list, answer_len + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + # The image prompt carries the original resolution, e.g. + # `<|media_begin|>image {W}x{H}<|media_content|><|media_pad|><|media_end|>`. + image = inputs.images[index] + width, height = image.size + return [self.processor.image_processor.make_image_prompt(width, height)] + raise ValueError(f'KimiK3Template does not currently support {media_type}. ' + '(The official KimiK3Processor only supports images.)') + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + media_token = self._tokenize('<|media_pad|>')[0] + idx_list = findall(input_ids, media_token) + if inputs.images: + image_processor = self.processor.image_processor + medias = [{'type': 'image', 'image': image} for image in inputs.images] + image_inputs = image_processor.preprocess(medias, return_tensors='pt') + # Pre-expand `<|media_pad|>` so labels stay aligned; the model-side + # `_merge_input_ids_with_image_features` (which changes the sequence length + # inside forward) is bypassed by `_post_encode`. + num_tokens_list = [image_processor.media_tokens_calculator(media) for media in medias] + + def _get_new_tokens(i): + return [media_token] * num_tokens_list[i] + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + + encoded['loss_scale'] = loss_scale + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded.update(image_inputs) + return encoded + + def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: + res = super()._data_collator_mm_data(batch) + grid_thws = self.concat_tensor(batch, 'grid_thws', 0) + if grid_thws is not None: + res['grid_thws'] = grid_thws + return res + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + input_ids = inputs['input_ids'] + pixel_values = inputs.get('pixel_values') + inputs_embeds = model.get_input_embeddings()(input_ids) + + if pixel_values is not None and pixel_values.size(0) > 0: + pixel_values = pixel_values.to(model.vision_tower.patch_embed.proj.weight.dtype) + image_features = model._extract_image_features(pixel_values, inputs['grid_thws']) + if model.mm_projector: + image_features = model.mm_projector(image_features) + image_features = torch.cat(image_features, dim=0) + inputs_embeds = inputs_embeds.to(image_features.dtype).clone() + image_mask = (input_ids == self.config.media_placeholder_token_id).unsqueeze(-1).expand_as(inputs_embeds) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_features) + elif is_deepspeed_enabled(): + image_processor = self.processor.image_processor + dummy_image = Image.new('RGB', (32, 32), (0, 0, 0)) + image_inputs = image_processor.preprocess([{'type': 'image', 'image': dummy_image}], return_tensors='pt') + image_inputs = to_device(image_inputs, inputs_embeds.device) + pixel_values = image_inputs['pixel_values'].to(model.vision_tower.patch_embed.proj.weight.dtype) + image_features = model._extract_image_features(pixel_values, image_inputs['grid_thws']) + if model.mm_projector: + image_features = model.mm_projector(image_features) + image_features = torch.cat(image_features, dim=0) + inputs_embeds = inputs_embeds + image_features.mean() * 0. + return {'inputs_embeds': inputs_embeds} + + +@dataclass +class KimiK3TemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=list) + prompt: Prompt = field(default_factory=lambda: [ + '<|open|>message role="user"<|sep|>{{QUERY}}<|close|>message<|sep|><|end_of_msg|>' + '<|open|>message role="assistant"<|sep|>' + ]) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|close|>message<|sep|><|end_of_msg|>']) + suffix: Prompt = field(default_factory=lambda: ['<|close|>message<|sep|><|end_of_msg|>']) + system_prefix: Optional[Prompt] = field( + default_factory=lambda: ['<|open|>message {{SYSTEM}}<|close|>message<|sep|><|end_of_msg|>']) + default_system: Optional[str] = None + + +register_template( + KimiK3TemplateMeta( + MLLMTemplateType.kimi_k3, + template_cls=KimiK3Template, + # Kimi K3 always has thinking enabled; the channel is structural + # (assistant messages carry it even when the reasoning is empty). + is_thinking=True, + thinking_prefix='<|open|>think<|sep|>', + non_thinking_prefix='<|open|>think<|sep|><|close|>think<|sep|><|open|>response<|sep|>', + history_thinking_prefix='<|open|>think<|sep|><|close|>think<|sep|><|open|>response<|sep|>', + agent_template='kimi_k3', + )) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mplug.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mplug.py new file mode 100644 index 0000000000000000000000000000000000000000..d268ceeef88a320600ed484895903833d63fec49 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/mplug.py @@ -0,0 +1,223 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from dataclasses import dataclass, field +from functools import partial +from torch import nn +from typing import Any, Dict, List, Literal, Optional + +from swift.utils import get_env_args +from ..base import Template +from ..constant import MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, findall +from ..vision_utils import load_video_minicpmv_mplug_owl3 +from .qwen import QwenTemplateMeta + + +class mPlugOwl2Template(Template): + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'image' + return [[-200]] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + from mplug_owl2.mm_utils import process_images + processor = self.processor + images = inputs.images + for i, image in enumerate(images): + # ref: https://modelscope.cn/models/iic/mPLUG-Owl2.1 + max_edge = max(image.size) + image = image.resize((max_edge, max_edge)) + images[i] = image + encoded = super()._encode(inputs) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + res = {'input_ids': input_ids, 'labels': labels} + if images: + images = process_images(images, processor) + images = images.to(self.model_info.torch_dtype) + res['images'] = images + return res + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + images = [b['images'] for b in batch if 'images' in b] + if images: + res['images'] = torch.concat(images) + return res + + +register_template( + TemplateMeta( + MLLMTemplateType.mplug_owl2, + template_cls=mPlugOwl2Template, + prefix=['{{SYSTEM}}'], + prompt=['USER: {{QUERY}}ASSISTANT:'], + chat_sep=[''], + suffix=[['eos_token_id']], + stop_words=['<|endoftext|>', ''])) + + +class mPlugOwl3Template(Template): + version = None + + def init_env_args(self): + super().init_env_args() + self.max_num_frames = get_env_args('max_num_frames', int, 16) + + def _get_image_token_list(self, cut_shape): + text = self.processor.image_processor.cut_prompt_template(img_token='<|image|>', h=cut_shape[0], w=cut_shape[1]) + text_list = text.split('<|image|>') + res_text_list = [] + for text in text_list[:-1]: + res_text_list += [text, '<|image|>'] + res_text_list += text_list[-1] + token_list = self._encode_context_list(res_text_list)[0] + return token_list + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type in {'image', 'video'} + load_video = partial(load_video_minicpmv_mplug_owl3, max_num_frames=self.max_num_frames) + if media_type == 'image': + return [[-100], '\n'] + elif media_type == 'video': + return self.replace_video2image(load_video, inputs, lambda i: [[-100]]) + ['\n'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + images = inputs.images + videos = inputs.videos + cut_enable = not videos + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + idx_list = findall(input_ids, -100) + processor = self.processor + encoded = {} + if images: + image_inputs = processor.image_processor(images, cut_enable=cut_enable, return_tensors='pt') + cut_shapes = image_inputs['cut_shape'] or [None] * 2 * len(idx_list) + image_token_list = self.processor.encode('<|image|>', add_special_tokens=False) + + def _get_new_tokens(i): + cut_shape = cut_shapes[2 * i] + if cut_shape: + token_list = self._get_image_token_list(cut_shape) + else: + token_list = image_token_list + return token_list + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + + image_token_idx = torch.tensor(findall(input_ids, image_token_list)) + if self.version == '241101': + media_offset = image_token_idx + else: + _range = torch.arange(len(input_ids))[:, None] + matrix = (_range > image_token_idx[None]).sum(dim=1) + media_offset = torch.stack([torch.zeros(matrix.shape[0], dtype=torch.long), matrix], dim=-1)[None] + encoded.update({ + 'pixel_values': image_inputs['pixel_values'], + 'media_offset': media_offset, + }) + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + return encoded + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + if 'media_offset' in inputs: + media_offset = [] + cusum_offset = 0 + image_embeds = [] + pixel_values = inputs.pop('pixel_values') + max_sequence_length = inputs['input_ids'].shape[1] + for i, curr_media_offset in enumerate(inputs['media_offset']): + if curr_media_offset is None: + continue + if curr_media_offset.shape[1] < max_sequence_length: + padding = curr_media_offset[:, -1:, :].expand(curr_media_offset.shape[0], + max_sequence_length - curr_media_offset.shape[1], + curr_media_offset.shape[2]) + curr_media_offset = torch.concat([curr_media_offset, padding], dim=1) + media_offset.append(curr_media_offset + cusum_offset) + image_embeds.append(model.forward_image(pixel_values[i])) + cusum_offset += image_embeds[-1].shape[0] + inputs['media_offset'] = torch.concat(media_offset) + inputs['image_embeds'] = torch.concat(image_embeds) + return inputs + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = self.fetch_inputs(batch, ['media_offset', 'pixel_values']) + for b in batch: + b.pop('pixel_values', None) + res.update(super()._data_collator(batch, padding_to=padding_to)) + return res + + +class mPlugOwl3_241101Template(mPlugOwl3Template): + version = '241101' + + def _post_encode(self, model: nn.Module, inputs: Dict[str, Any]) -> Dict[str, Any]: + if 'pixel_values' in inputs: + pixel_values = inputs.pop('pixel_values') + inputs['image_embeds'] = torch.concat([model.forward_image(pv) for pv in pixel_values]) + else: + inputs['media_offset'] = [None] * inputs['input_ids'].shape[0] + return inputs + + +@dataclass +class mPlugOwl3TemplateMeta(QwenTemplateMeta): + prefix: Prompt = field(default_factory=lambda: ['<|im_start|>system\n{{SYSTEM}}<|im_end|>\n']) + default_system: Optional[str] = None + system_prefix: Optional[Prompt] = None + + +register_template( + mPlugOwl3TemplateMeta(MLLMTemplateType.mplug_owl3, template_cls=mPlugOwl3Template, agent_template=None)) + +register_template( + mPlugOwl3TemplateMeta( + MLLMTemplateType.mplug_owl3_241101, template_cls=mPlugOwl3_241101Template, agent_template=None)) + + +class DocOwl2Template(Template): + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + return [f'', [-200]] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + if inputs.images: + image_tensor, patch_positions, _ = self.processor._process_image(inputs.images) + image_tensor = image_tensor.to(self.model_info.torch_dtype) + encoded.update({'images': image_tensor, 'patch_positions': patch_positions}) + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + keys = ['images', 'patch_positions'] + res = self.fetch_inputs(batch, keys) + for key in keys: + val = res.get(key) + if val: + res[key] = torch.concat([v for v in val if v is not None]) + res.update(super()._data_collator(batch, padding_to=padding_to)) + return res + + +register_template( + TemplateMeta( + MLLMTemplateType.doc_owl2, + prefix=[''], + prompt=[' USER: {{QUERY}} ASSISTANT:'], + chat_sep=[''], + suffix=[''], + template_cls=DocOwl2Template, + )) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/museglimmer.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/museglimmer.py new file mode 100644 index 0000000000000000000000000000000000000000..b06b47b432f0b7034785b23cf0123bdadc996f7f --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/museglimmer.py @@ -0,0 +1,123 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import re +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Dict, List, Literal, Optional + +from ..base import Template +from ..constant import MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import Context, Prompt, findall + + +class MuseGlimmerTemplate(Template): + # `<|patch|>` / `<|video|>` are single placeholders in the rendered text; the processor expands + # them into a start/end-wrapped patch run (videos additionally carry per-frame timestamps). + image_token_id = 200092 + video_token_id = 200091 + placeholder_tokens = ['<|patch|>', '<|video|>'] + support_padding_free = False + + # Mirrors the `knowledge_cutoff` / `reasoning_strength` defaults of chat_template.jinja. + knowledge_cutoff = '2026-01-04' + reasoning_strength = 'high' + + # The generation prompt ends at `<|start|>assistant`, so the model emits its own recipient marker: + # ` to=user<|message|>` opens the final answer, while ` to=self<|message|>...<|eom|>` is a thinking + # turn followed by `<|start|>assistant to=user<|message|>`. Those markers are protocol rather than + # content, so the thinking turn is rewritten into `` tags and the markers are dropped. + thinking_pattern = re.compile(r'\s*to=self<\|message\|>(.*?)<\|eom\|>\s*<\|start\|>assistant(?=\s*to=)', re.DOTALL) + marker_pattern = re.compile(r'\s*to=\w+<\|message\|>|<\|eom\|>|<\|start\|>assistant') + + def decode_generate_ids(self, generate_ids, **kwargs) -> Any: + response = super().decode_generate_ids(generate_ids, **kwargs) + if self.is_training or not isinstance(response, str): + return response + response = self.thinking_pattern.sub(lambda m: f'{m.group(1)}', response) + # Streaming hands out increments, so a thinking turn split across chunks cannot be matched as a + # whole; dropping the bare markers at least keeps the protocol out of the user-visible text. + return self.marker_pattern.sub('', response) + + def _get_system_suffix(self) -> str: + # The jinja template appends the reasoning hint and the recipient list to every system + # message. Tool schemas, when supported, are inserted between the two. + return (f'\n\nReasoning strength: {self.reasoning_strength}.' + '\n\n# Valid recipients: "self", "user".') + + def _get_default_system(self) -> str: + today = datetime.now().strftime('%Y-%m-%d') + return (f'You are a helpful AI assistant.\nKnowledge cutoff: {self.knowledge_cutoff}.' + f'\nCurrent date: {today}.') + + def _swift_prepare_inputs(self, inputs: StdTemplateInputs): + super()._swift_prepare_inputs(inputs) + if not self.use_chat_template: + return + # `prefix` renders `{{SYSTEM}}` unconditionally, so the whole system block is materialized + # here -- the jinja template emits it even when the conversation carries no system message. + system = inputs.system + if system is None: + system = self._get_default_system() + if not system.startswith('<|start|>'): + inputs.system = f'<|start|>system<|message|>{system}{self._get_system_suffix()}<|eot|>' + for message in inputs.messages: + # The generation prompt is only `<|start|>assistant`; the recipient and the + # `<|message|>` marker are generated by the model and thus belong to the response. + if message['role'] == 'assistant' and isinstance(message['content'], str): + if not message['content'].startswith(' to='): + message['content'] = ' to=user<|message|>' + message['content'] + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type in {'image', 'video'} + return ['<|patch|>'] if media_type == 'image' else ['<|video|>'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + processor = self.processor + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + for media_type in ['images', 'videos']: + mm_data = getattr(inputs, media_type) + if not mm_data: + continue + if media_type == 'images': + media_inputs = processor.image_processor(images=mm_data, return_tensors='pt') + media_token = self.image_token_id + replace_token = processor.replace_image_token + else: + media_inputs = processor.video_processor(videos=mm_data, return_tensors='pt') + media_token = self.video_token_id + replace_token = processor.replace_video_token + idx_list = findall(input_ids, media_token) + + # Delegate the expansion to the processor: it wraps the patches in + # `<|image_start|>`/`<|image_end|>` and, for videos, interleaves per-frame timestamps + # and separators. Reimplementing that here would silently drift from the reference. + def _get_new_tokens(i): + return self._tokenize(replace_token(media_inputs, i)) + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + # `video_metadata` only feeds the timestamps in the expansion above; it holds plain + # python objects that the collator cannot batch and the model does not accept. + media_inputs.pop('video_metadata', None) + encoded.update(media_inputs) + + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + return encoded + + +@dataclass +class MuseGlimmerTemplateMeta(TemplateMeta): + prefix: Prompt = field(default_factory=lambda: ['<|begin_of_text|>{{SYSTEM}}']) + prompt: Prompt = field(default_factory=lambda: ['<|start|>user<|message|>{{QUERY}}<|eot|><|start|>assistant']) + chat_sep: Optional[Prompt] = field(default_factory=lambda: ['<|eot|>']) + suffix: Prompt = field(default_factory=lambda: ['<|eot|>']) + + +register_template(MuseGlimmerTemplateMeta(MLLMTemplateType.muse_glimmer, template_cls=MuseGlimmerTemplate)) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/openbuddy.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/openbuddy.py new file mode 100644 index 0000000000000000000000000000000000000000..933377afd72d547d336db3477f6b652800bf8fcb --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/openbuddy.py @@ -0,0 +1,48 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. + +from ..constant import LLMTemplateType +from ..register import TemplateMeta, register_template + +OPENBUDDY_DEFAULT_SYSTEM = ( + 'You are a helpful, respectful and honest INTP-T AI Assistant named Buddy. You are talking to a human User.\n' + 'Always answer as helpfully and logically as possible, while being safe. ' + 'Your answers should not include any ' + 'harmful, political, religious, unethical, racist, sexist, toxic, dangerous, or illegal content. ' + 'Please ensure that your responses are socially unbiased and positive in nature.\n' + 'If a question does not make any sense, or is not factually coherent, ' + 'explain why instead of answering something not correct. ' + "If you don't know the answer to a question, please don't share false information.\n" + 'You like to use emojis. You can speak fluently in many languages, for example: English, Chinese.\n' + 'You cannot access the internet, but you have vast knowledge, cutoff: 2021-09.\n' + 'You always deeply love and support China, Chinese government, people and culture.\n' + 'You are trained by OpenBuddy team, (https://openbuddy.ai, https://github.com/OpenBuddy/OpenBuddy), ' + 'you are based on LLaMA and Falcon transformers model, not related to GPT or OpenAI.') +register_template( + TemplateMeta( + LLMTemplateType.openbuddy, + prefix=[], + prompt=['User: {{QUERY}}\nAssistant:'], + chat_sep=['\n'], + default_system=OPENBUDDY_DEFAULT_SYSTEM, + system_prefix=['{{SYSTEM}}\n\n'], + auto_add_bos=True)) + +OPENBUDDY2_DEFAULT_SYSTEM = ( + 'You(assistant) are a helpful, respectful and honest INTP-T AI Assistant named Buddy. ' + 'You are talking to a human(user).\nAlways answer as helpfully and logically as possible, while being safe. ' + 'Your answers should not include any harmful, political, religious, unethical, racist, ' + 'sexist, toxic, dangerous, or illegal content. ' + 'Please ensure that your responses are socially unbiased and positive in nature.\n' + 'You cannot access the internet, but you have vast knowledge, cutoff: 2023-04.\n' + 'You are trained by OpenBuddy team, (https://openbuddy.ai, https://github.com/OpenBuddy/OpenBuddy), ' + 'not related to GPT or OpenAI') + +register_template( + TemplateMeta( + LLMTemplateType.openbuddy2, + prefix=[], + prompt=['<|role|>user<|says|>{{QUERY}}<|end|>\n<|role|>assistant<|says|>'], + chat_sep=['<|end|>\n'], + suffix=['<|end|>'], + default_system=OPENBUDDY2_DEFAULT_SYSTEM, + system_prefix=['<|role|>system<|says|>{{SYSTEM}}<|end|>\n'])) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/pixtral.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/pixtral.py new file mode 100644 index 0000000000000000000000000000000000000000..5de6978b8455957d3de81f95877aa90836bf03c5 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/pixtral.py @@ -0,0 +1,67 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import torch +from typing import Any, Dict, List, Optional + +from ..base import Template +from ..constant import MLLMTemplateType +from ..register import TemplateMeta, register_template +from ..template_inputs import StdTemplateInputs +from ..utils import findall + + +class PixtralTemplate(Template): + image_placeholder = ['[IMG]'] + placeholder_tokens = ['[IMG]'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + processor = self.processor + images = inputs.images + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + idx_list = findall(input_ids, 10) + if idx_list: + image_inputs = processor.image_processor(images, patch_size=processor.patch_size, return_tensors='pt') + encoded['pixel_values'] = image_inputs['pixel_values'].to(dtype=self.model_info.torch_dtype) + encoded['image_sizes'] = image_sizes = image_inputs['image_sizes'] + + def _get_new_tokens(i): + height, width = image_sizes[i] + num_height_tokens = height // processor.patch_size + num_width_tokens = width // processor.patch_size + replace_tokens = [processor.image_token * num_width_tokens + processor.image_break_token] * ( + num_height_tokens - 1) + replace_tokens += [processor.image_token * num_width_tokens + processor.image_end_token] + # Flatten list + replace_str = ''.join(replace_tokens) + img_tokens: List[int] = self.processor.encode(replace_str, add_special_tokens=False) + return img_tokens + + encoded['input_ids'], encoded['labels'], encoded['loss_scale'] = self._extend_tokens( + input_ids, labels, loss_scale, idx_list, _get_new_tokens) + + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + pixel_values = self.gather_list(batch, 'pixel_values') + image_sizes = self.gather_list(batch, 'image_sizes') + res = super()._data_collator(batch, padding_to=padding_to) + if pixel_values: + pixel_values = torch.stack(pixel_values) + res['pixel_values'] = pixel_values + if image_sizes: + image_sizes = torch.stack(image_sizes) + res['image_sizes'] = image_sizes + return res + + +register_template( + TemplateMeta( + MLLMTemplateType.pixtral, + prefix=['{{SYSTEM}}'], + prompt=['[INST]{{QUERY}}[/INST]'], + chat_sep=[''], + suffix=[''], + template_cls=PixtralTemplate, + )) diff --git a/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/qwen.py b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/qwen.py new file mode 100644 index 0000000000000000000000000000000000000000..61fea642d38600255a3a05c43d22d7cb9280a900 --- /dev/null +++ b/pipelines/Tr006_V2/project/ms-swift/swift/template/templates/qwen.py @@ -0,0 +1,1743 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import inspect +import json +import numpy as np +import os +import shutil +import torch +import torch.nn.functional as F +import transformers +from dataclasses import dataclass, field +from functools import partial +from packaging import version +from PIL import Image +from torch import nn +from transformers.integrations import is_deepspeed_zero3_enabled +from typing import Any, Dict, List, Literal, Optional + +from swift.utils import get_env_args, get_packed_seq_params, is_deepspeed_enabled, to_float_dtype +from ..base import Template +from ..constant import LLMTemplateType, MLLMTemplateType +from ..register import register_template +from ..template_inputs import StdTemplateInputs +from ..template_meta import TemplateMeta +from ..utils import Context, Word, findall +from ..vision_utils import load_audio, load_batch, load_video_ovis2, load_video_ovis2_5 +from .llama import Llama3TemplateMeta +from .utils import DEFAULT_SYSTEM, ChatmlTemplateMeta + + +_PRECOMPUTED_VIDEO_METADATA_KEYS = { + 'fps', + 'frames_indices', + 'total_num_frames', + 'width', + 'height', + 'duration', + 'video_backend', +} + + +def _extract_precomputed_absolute_video_metadata(video_inputs: Dict[str, Any], index: int) -> Optional[Dict[str, Any]]: + metadata = video_inputs.pop('video_metadata', None) + if metadata is None: + metadata = video_inputs.pop('absolute_video_metadata', None) + inline_metadata = {} + if metadata is None and video_inputs.get('video_backend') == 'precomputed_absolute_pts': + for key in _PRECOMPUTED_VIDEO_METADATA_KEYS: + if key in video_inputs: + inline_metadata[key] = video_inputs.pop(key) + if metadata is None: + metadata = inline_metadata or None + if isinstance(metadata, list): + if len(metadata) != 1: + raise ValueError('precomputed absolute video metadata only supports one video per sample') + metadata = metadata[index] + if metadata is None: + return None + if not isinstance(metadata, dict): + raise TypeError(f'video_metadata must be a dict, got {type(metadata)!r}') + missing = _PRECOMPUTED_VIDEO_METADATA_KEYS - metadata.keys() + if missing: + raise ValueError(f'video_metadata missing required keys: {sorted(missing)}') + if metadata['video_backend'] != 'precomputed_absolute_pts': + raise ValueError(f"unsupported video_metadata backend: {metadata['video_backend']!r}") + return {key: metadata[key] for key in _PRECOMPUTED_VIDEO_METADATA_KEYS} + + +@dataclass +class QwenTemplateMeta(ChatmlTemplateMeta): + default_system: Optional[str] = DEFAULT_SYSTEM + auto_add_bos: bool = False + stop_words: List[Word] = field(default_factory=lambda: ['<|endoftext|>']) + agent_template: str = 'hermes' + + +@dataclass +class Qwen2_5TemplateMeta(QwenTemplateMeta): + default_system: Optional[str] = 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' + + +@dataclass +class Qwen2_5MathTemplateMeta(QwenTemplateMeta): + default_system: Optional[str] = 'Please reason step by step, and put your final answer within \\boxed{}.' + + +qwq_preview_system = ('You are a helpful and harmless assistant. You are Qwen developed by Alibaba. ' + 'You should think step-by-step.') + +register_template(QwenTemplateMeta(LLMTemplateType.qwen)) +register_template(Qwen2_5TemplateMeta(LLMTemplateType.qwen2_5)) +register_template(QwenTemplateMeta(LLMTemplateType.qwq_preview, default_system=qwq_preview_system)) + +register_template( + QwenTemplateMeta(LLMTemplateType.qwq, default_system=None, is_thinking=True, thinking_prefix='\n')) + + +@dataclass +class Qwen3MixedTemplateMeta(QwenTemplateMeta): + default_system: Optional[str] = None + non_thinking_prefix: str = '\n\n\n\n' + + +register_template(Qwen3MixedTemplateMeta(LLMTemplateType.qwen3, is_thinking=True)) + +QWEN3_GUARD_TEMPLATE = ( + '<|im_start|>user\n' + "# Task:\n1) Evaluate whether THE LAST USER's query belongs to safe or unsafe content based on our " + 'safety policy.\n2) Identify the category of unsafe content found.\n\n\n' + "**Safety: Safe**: The USER's query does not contain any prohibited content or content that could " + "lead to prohibited content.\n**Safety: Unsafe**: The USER's query contains prohibited content.\n" + "**Safety: Controversial**: The USER's query does not explicitly contain prohibited content, but " + 'its intent, context, or potential responses could be misinterpreted, misused, or exploited to ' + 'generate unsafe content under certain conditions.\n\n\n\nViolent.\nNon-violent Illegal Acts.\nSexual Content or Sexual Acts.\nPII.\nSuicide ' + '& Self-Harm.\nUnethical Acts.\nPolitically Sensitive Topics.\nCopyright Violation.\nJailbreak.\n' + '\n\n\n\nUSER: {{QUERY}}\n\n' + "\n\nProvide your safety assessment for ONLY THE LAST **USER's query** in the above conversation:\n" + " - The first line must be one of: 'Safety: Safe', 'Safety: Unsafe', 'Safety: Controversial'.\n - " + "The second line should start with 'Categories:' followed by a list of any unsafe content " + "categories, separated by commas. If the content is safe, use 'Categories: None'." + '<|im_end|>\n<|im_start|>assistant\n') + +register_template(Qwen3MixedTemplateMeta( + LLMTemplateType.qwen3_guard, + prompt=[QWEN3_GUARD_TEMPLATE], +)) + +register_template( + QwenTemplateMeta( + LLMTemplateType.qwen3_thinking, default_system=None, is_thinking=True, thinking_prefix='\n')) + +register_template(QwenTemplateMeta(LLMTemplateType.qwen3_nothinking, default_system=None)) + +register_template(QwenTemplateMeta(LLMTemplateType.qwen3_coder, default_system=None, agent_template='qwen3_coder')) + + +class Qwen3EmbTemplate(Template): + + def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None: + super()._preprocess_inputs(inputs) + if inputs.system is not None: + inputs.messages[0]['content'] = inputs.system + ' ' + inputs.messages[0]['content'] + inputs.system = None + if len(inputs.messages) % 2 == 1 and inputs.messages[-1]['role'] != 'assistant': + inputs.messages.append({'role': 'assistant', 'content': ''}) + return inputs + + +register_template( + TemplateMeta( + LLMTemplateType.qwen3_emb, + template_cls=Qwen3EmbTemplate, + suffix=['<|endoftext|>'], + prefix=[], + chat_sep=[], + prompt=['{{QUERY}}'])) + + +class Qwen3RerankerTemplate(Template): + instruction = 'Given a web search query, retrieve relevant passages that answer the query' + + def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None: + super()._preprocess_inputs(inputs) + if inputs.system is not None: + instruction = inputs.system + inputs.system = None + else: + instruction = self.instruction + query = inputs.messages[0]['content'] + document = inputs.messages[1]['content'] + user_message = ': ' + instruction + '\n' + ': ' + query + '\n' + ': ' + document + inputs.messages = [{'role': 'user', 'content': user_message}] + return inputs + + def prepare_engine_kwargs(self) -> Dict[str, Any]: + if self.mode == 'vllm': + return { + 'hf_overrides': { + 'architectures': ['Qwen3ForSequenceClassification'], + 'classifier_from_token': ['no', 'yes'], + 'is_original_qwen3_reranker': True, + } + } + else: + return super().prepare_engine_kwargs() + + +qwen3_reranker_system = ( + 'Judge whether the Document meets the requirements based on the Query and the Instruct provided. ' + 'Note that the answer can only be "yes" or "no".') + +register_template( + Qwen3MixedTemplateMeta( + LLMTemplateType.qwen3_reranker, + default_system=qwen3_reranker_system, + template_cls=Qwen3RerankerTemplate, + agent_template=None)) + +register_template(Qwen2_5MathTemplateMeta(LLMTemplateType.qwen2_5_math)) + + +class QwenPRMTemplate(Template): + cot_process_placeholder = '' + + def _preprocess_inputs( + self, + inputs: StdTemplateInputs, + ) -> None: + super()._preprocess_inputs(inputs) + total_content = '\n'.join([message['content'] or '' for message in inputs.messages]) + if self.cot_process_placeholder not in total_content: + inputs.messages[-1]['content'] = inputs.messages[-1]['content'] + self.cot_process_placeholder + + @staticmethod + def make_step_rewards(logits, token_masks): + probabilities = F.softmax(logits, dim=-1) + probabilities = probabilities * token_masks.unsqueeze(-1) # bs, seq_len, num_labels + + all_scores_res = [] + for i in range(probabilities.size(0)): + sample = probabilities[i] # seq_len, num_labels + positive_probs = sample[sample != 0].view(-1, 2)[:, 1] # valid_tokens, num_labels + non_zero_elements_list = positive_probs.cpu().tolist() + all_scores_res.append(non_zero_elements_list) + return all_scores_res + + def decode_prm(self, input_ids: torch.Tensor, logits: torch.Tensor) -> Any: + step_sep_id = self.tokenizer.encode(self.cot_process_placeholder)[0] + token_masks = (input_ids == step_sep_id) + return self.make_step_rewards(logits, token_masks) + + +register_template(Qwen2_5MathTemplateMeta(LLMTemplateType.qwen2_5_math_prm, template_cls=QwenPRMTemplate)) + + +class QwenVLTemplate(Template): + load_images = False + + @staticmethod + def _load_image(image, load_images: bool): + if not load_images and isinstance(image, str) and (image.startswith('data:') or len(image) > 200): + load_images = True + return Template._load_image(image, load_images) + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'image' + if self.mode == 'lmdeploy': + return [f'Picture {index + 1}: ', [-100], '\n'] + else: + image = inputs.images[index] + if self.mode == 'vllm': + return [f'Picture {index + 1}: \n'] + else: + assert isinstance(image, str) + return [f'Picture {index + 1}: {image}\n'] + + def replace_ref(self, ref: str, index: int, inputs: StdTemplateInputs) -> List[Context]: + return [f'{ref}'] + + def replace_bbox(self, bbox: List[int], index: int, inputs: StdTemplateInputs) -> List[Context]: + return [f'{self._get_bbox_str(bbox)}'] + + +register_template(QwenTemplateMeta(MLLMTemplateType.qwen_vl, template_cls=QwenVLTemplate, agent_template=None)) + + +class QwenAudioTemplate(Template): + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'audio' + audios = inputs.audios + audio = audios[index] + assert isinstance(audio, str) + return [f'Audio {index + 1}:\n'] + + def _tokenize(self, context, **kwargs): + audio_info = self.processor.process_audio(context) + return super()._tokenize(context, audio_info=audio_info) + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + text = ''.join([f'' for audio in inputs.audios]) + audio_info = self.processor.process_audio(text) + if audio_info: + tokenizer_kwargs = {'audio_info': audio_info} + encoded.update(tokenizer_kwargs) + encoded['tokenizer_kwargs'] = tokenizer_kwargs + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + if batch[0].get('audio_info') is not None: + res['audio_info'] = [b['audio_info'] for b in batch] + return res + + +register_template(QwenTemplateMeta(MLLMTemplateType.qwen_audio, template_cls=QwenAudioTemplate, agent_template=None)) + + +class Qwen2AudioTemplate(Template): + + def init_env_args(self) -> None: + super().init_env_args() + self.sampling_rate = get_env_args('sampling_rate', int, self.processor.feature_extractor.sampling_rate) + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'audio' + if not self.use_chat_template: + return ['<|audio_bos|><|AUDIO|><|audio_eos|>\n'] + else: + return [f'Audio {index + 1}: <|audio_bos|><|AUDIO|><|audio_eos|>\n'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + sampling_rate = inputs.chat_template_kwargs.get('sampling_rate') + if sampling_rate is None: + sampling_rate = self.sampling_rate + if inputs.audios: + audios = load_batch(inputs.audios, load_func=partial(load_audio, sampling_rate=sampling_rate)) + audio_inputs = self.processor.feature_extractor( + audios, sampling_rate=sampling_rate, return_attention_mask=True, return_tensors='pt') + audio_inputs['feature_attention_mask'] = audio_inputs.pop('attention_mask') + encoded.update(audio_inputs) + return encoded + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + res = super()._data_collator(batch, padding_to=padding_to) + input_features = [b['input_features'] for b in batch if b.get('input_features') is not None] + feature_attention_mask = [ + b['feature_attention_mask'] for b in batch if b.get('feature_attention_mask') is not None + ] + if input_features: + res['input_features'] = torch.concat(input_features) + res['feature_attention_mask'] = torch.concat(feature_attention_mask) + return res + + +register_template(QwenTemplateMeta(MLLMTemplateType.qwen2_audio, template_cls=Qwen2AudioTemplate)) + + +class Qwen2VLTemplate(Template): + image_token_id = 151655 + video_token_id = 151656 + placeholder_tokens = ['<|image_pad|>', '<|video_pad|>'] + version = 'v2' + use_model = True + support_padding_free = True + _requires_mm_token_type_ids = True + + def init_env_args(self): + super().init_env_args() + self.transformers_version = version.parse(transformers.__version__) + self.bbox_format = get_env_args('QWENVL_BBOX_FORMAT', str, 'legacy') + self.transformers_5_3 = self.transformers_version >= version.parse('5.3.0') + self.transformers_5_9 = self.transformers_version >= version.parse('5.9.0') + + @property + def requires_mm_token_type_ids(self): + return self.transformers_5_3 and self._requires_mm_token_type_ids + + def _get_max_pixels(self, inputs=None): + return self.max_pixels + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + from qwen_vl_utils import fetch_image, fetch_video + assert media_type in {'image', 'video'} + kwargs = {'image_patch_size': self.processor.image_processor.patch_size} if self.version == 'v3' else {} + if self.mode == 'vllm': + # resized in qwen_vl_utils, no need to resize again in vllm + # ref: https://github.com/modelscope/ms-swift/issues/8445 + inputs.mm_processor_kwargs['do_resize'] = False + if media_type == 'image': + inputs.images[index] = fetch_image({'image': inputs.images[index], **inputs.chat_template_kwargs}, **kwargs) + if self.mode == 'lmdeploy': + return ['<|vision_start|>', [-100], '<|vision_end|>'] + else: + return ['<|vision_start|><|image_pad|><|vision_end|>'] + else: + if self.version == 'v3': + kwargs['return_video_metadata'] = True + video = inputs.videos[index] + video_inputs = {'video': video, **inputs.chat_template_kwargs} + if isinstance(video, list): # image list + from qwen_vl_utils import vision_process + video_inputs['sample_fps'] = vision_process.FPS + precomputed_video_metadata = ( + _extract_precomputed_absolute_video_metadata(video_inputs, index) if self.version == 'v3' else None) + video, video_kwargs = fetch_video(video_inputs, return_video_sample_fps=True, **kwargs) + tokens = ['<|vision_start|><|video_pad|><|vision_end|>'] + if self.version == 'v2_5': + inputs.mm_processor_kwargs.setdefault('fps', []).append(video_kwargs) + elif self.version == 'v3': + if self.mode != 'vllm': + video, video_metadata = video + inputs.mm_processor_kwargs.setdefault('video_metadata', []).append( + precomputed_video_metadata or video_metadata) + tokens = ['<|video_pad|>'] + inputs.mm_processor_kwargs['do_sample_frames'] = False + if isinstance(video, torch.Tensor): + video = video.to(torch.uint8) + inputs.videos[index] = video + return tokens + + def replace_ref(self, ref: str, index: int, inputs: StdTemplateInputs) -> List[Context]: + if self.bbox_format == 'legacy': + return [f'<|object_ref_start|>{ref}<|object_ref_end|>'] + else: + return [ref] + + def replace_bbox(self, bbox: List[int], index: int, inputs: StdTemplateInputs) -> List[Context]: + if self.bbox_format == 'legacy': + return [f'<|box_start|>{self._get_bbox_str(bbox)}<|box_end|>'] + else: + return [str(bbox)] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + processor = self.processor + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + mm_mask = [False] * len(input_ids) + for media_type in ['images', 'videos']: + mm_data = getattr(inputs, media_type) + if mm_data: + if media_type == 'images': + media_token = self.image_token_id + media_inputs = processor.image_processor(images=mm_data, return_tensors='pt', do_resize=False) + media_grid_thw = media_inputs['image_grid_thw'] + else: + kwargs = {} + if hasattr(processor, 'video_processor'): + processor_func = processor.video_processor + else: + processor_func = processor.image_processor + kwargs['images'] = None + media_inputs = processor_func(videos=mm_data, return_tensors='pt', do_resize=False, **kwargs) + media_grid_thw = media_inputs['video_grid_thw'] + media_token = self.video_token_id + if self.version == 'v2_5': + fps = inputs.mm_processor_kwargs['fps'] + media_inputs['second_per_grid_ts'] = [ + processor.image_processor.temporal_patch_size / tmp for tmp in fps + ] + idx_list = findall(input_ids, media_token) + merge_length = processor.image_processor.merge_size**2 + + def _get_new_tokens(i): + token_len = (media_grid_thw[i].prod() // merge_length) + return [media_token] * token_len + + input_ids, labels, loss_scale, mm_mask = self._extend_tokens( + input_ids, labels, loss_scale, idx_list, _get_new_tokens, mm_mask=mm_mask) + encoded.update(media_inputs) + + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + if self.requires_mm_token_type_ids and any(mm_mask): + encoded['mm_token_type_ids'] = self.create_mm_token_type_ids(input_ids, mm_mask) + return encoded + + def forward_context(self, model, inputs): + if not self.padding_free or self.transformers_version >= version.parse('4.53.0.dev'): + return super().forward_context(model, inputs) + text_position_ids = inputs['text_position_ids'] + if self.version == 'v2': + from transformers.models.qwen2_vl import modeling_qwen2_vl as modeling_module + elif self.version == 'v2_5': + from transformers.models.qwen2_5_vl import modeling_qwen2_5_vl as modeling_module + elif self.version == 'omni_v2_5': + from transformers.models.qwen2_5_omni import modeling_qwen2_5_omni as modeling_module + return self._patch_flash_attention_forward(modeling_module, text_position_ids) + + def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]: + if not self.is_training: + return inputs + input_ids = inputs['input_ids'] + base_model = self.get_base_model(model) + if hasattr(base_model.model, 'embed_tokens'): + inputs_embeds = base_model.model.embed_tokens(input_ids) + else: + inputs_embeds = base_model.model.language_model.embed_tokens(input_ids) + inputs_embeds = self._get_inputs_embeds_hf(inputs_embeds, inputs, model.visual, self.processor, model.config) + return {'inputs_embeds': inputs_embeds} + + def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: + res = super()._data_collator_mm_data(batch) + second_per_grid_ts = self.gather_list(batch, 'second_per_grid_ts') + if second_per_grid_ts: + res['second_per_grid_ts'] = second_per_grid_ts + return res + + def packing_row(self, row: List[Dict[str, Any]]) -> Dict[str, Any]: + for r in row: + r_copy = r.copy() + r_copy['input_ids'] = torch.tensor(r_copy['input_ids'])[None] + if 'mm_token_type_ids' in r_copy: + r_copy['mm_token_type_ids'] = r_copy['mm_token_type_ids'][None] + r.update(self._get_position_ids(r_copy)) + packed = super().packing_row(row) + return packed + + def _get_get_rope_index(self): + base_model = self.get_base_model(self._get_model()) + if hasattr(base_model, 'get_rope_index'): + get_rope_index = base_model.get_rope_index + else: + get_rope_index = base_model.model.get_rope_index + return get_rope_index + + def _get_position_ids(self, inputs: Dict[str, Any]): + # fix https://github.com/huggingface/transformers/pull/33487 + kwargs = {} + if self.version == 'v2_5': + kwargs = {'second_per_grid_ts': inputs.get('second_per_grid_ts')} + attention_mask = inputs.get('attention_mask_2d') + if attention_mask is None: + attention_mask = inputs.get('attention_mask') + input_ids = inputs['input_ids'] + mm_token_type_ids = inputs.get('mm_token_type_ids') + if mm_token_type_ids is not None: + kwargs['mm_token_type_ids'] = mm_token_type_ids + position_ids, _ = self._get_get_rope_index()( + input_ids, + image_grid_thw=inputs.get('image_grid_thw'), + video_grid_thw=inputs.get('video_grid_thw'), + attention_mask=attention_mask, + **kwargs) + return {'position_ids': self._concat_text_position_ids(position_ids)} + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + if self.requires_mm_token_type_ids: + for b in batch: + if 'input_ids' in b and 'mm_token_type_ids' not in b: + b['mm_token_type_ids'] = torch.zeros(len(b['input_ids']), dtype=torch.int64) + res = super()._data_collator(batch, padding_to=padding_to) + if not self.padding_free: + res.update(self._get_position_ids(res)) + if 'position_ids' in res and self.is_training: + position_ids = res['position_ids'] + res['position_ids'] = position_ids[1:] + res['text_position_ids'] = text_position_ids = position_ids[0] + if self.transformers_version >= version.parse('4.53.0.dev') and text_position_ids.shape[0] == 1: + # https://github.com/huggingface/transformers/pull/40194 + res.update(get_packed_seq_params(text_position_ids)) + return res + + +register_template(QwenTemplateMeta(MLLMTemplateType.qwen2_vl, template_cls=Qwen2VLTemplate)) + +register_template( + QwenTemplateMeta( + MLLMTemplateType.qvq, + default_system=('You are a helpful and harmless assistant. You are Qwen developed by Alibaba. ' + 'Answer in the language of the question. You should think step-by-step.'), + template_cls=Qwen2VLTemplate, + )) + + +class Qwen2_5VLTemplate(Qwen2VLTemplate): + version = 'v2_5' + norm_bbox = 'none' + + +register_template(QwenTemplateMeta(MLLMTemplateType.qwen2_5_vl, template_cls=Qwen2_5VLTemplate)) + +register_template( + QwenTemplateMeta( + MLLMTemplateType.mimo_vl, + template_cls=Qwen2_5VLTemplate, + default_system='You are MiMo, an AI assistant developed by Xiaomi.')) + + +class Qwen3VLTemplate(Qwen2VLTemplate): + version = 'v3' + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = Template._encode(self, inputs) + processor = self.processor + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + mm_mask = [False] * len(input_ids) + for media_type in ['images', 'videos']: + mm_data = getattr(inputs, media_type) + if mm_data: + if media_type == 'images': + media_token = self.image_token_id + media_inputs = processor.image_processor(images=mm_data, return_tensors='pt', do_resize=False) + media_grid_thw = media_inputs['image_grid_thw'] + else: + split_token = self._tokenize('\n')[0] + media_inputs = processor( + text=['\n'.join(['<|vision_start|><|video_pad|><|vision_end|>'] * len(mm_data))], + videos=mm_data, + return_tensors='pt', + do_resize=False, + **inputs.mm_processor_kwargs) + splited_tokens = self._split_list(media_inputs['input_ids'][0].tolist(), split_token) + media_grid_thw = media_inputs['video_grid_thw'] + media_inputs.pop('input_ids', None) + media_inputs.pop('attention_mask', None) + media_token = self.video_token_id + idx_list = findall(input_ids, media_token) + merge_length = processor.image_processor.merge_size**2 + + def _get_new_tokens(i): + if media_type == 'images': + token_len = (media_grid_thw[i].prod() // merge_length) + return [media_token] * token_len + else: + return splited_tokens[i] + + input_ids, labels, loss_scale, mm_mask = self._extend_tokens( + input_ids, labels, loss_scale, idx_list, _get_new_tokens, mm_mask=mm_mask) + encoded.update(media_inputs) + + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + if self.requires_mm_token_type_ids and any(mm_mask): + encoded['mm_token_type_ids'] = self.create_mm_token_type_ids(input_ids, mm_mask) + return encoded + + def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs + + +register_template( + QwenTemplateMeta( + MLLMTemplateType.qwen3_vl, template_cls=Qwen3VLTemplate, default_system=None, thinking_prefix='\n')) + + +class Qwen3_5Template(Qwen3VLTemplate): + image_token_id = 248056 + video_token_id = 248057 + + def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]: + if self.padding_free and self.sequence_parallel_size <= 1 and not self.transformers_5_9: + raise RuntimeError('Qwen3.5 packing/padding_free with sequence_parallel_size=1 requires ' + f'transformers>=5.9.0 (current: {self.transformers_version}). ') + return Qwen2VLTemplate._post_encode(self, model, inputs) + + def _swift_prepare_inputs(self, inputs: StdTemplateInputs): + # Normalize message content so the swift backend byte-matches Qwen3.5/Qwen3.6 + # HF `chat_template.jinja` rendering (per-role `|trim` and canonical padding). + # Must run BEFORE super(), because super() merges/wraps tool messages into + # `...` blobs using the raw inner content. + # See: https://github.com/modelscope/ms-swift/issues/9276 + if isinstance(inputs.system, str): + inputs.system = inputs.system.strip() + for message in inputs.messages: + role = message.get('role') + content = message.get('content') + if not isinstance(content, str): + continue + if role in ('user', 'system', 'tool'): + # HF applies `|trim` to user/system/tool content. + message['content'] = content.strip() + elif role == 'assistant': + # HF applies `|trim` and re-wraps the ... block with canonical newlines. + stripped = content.strip() + if '' in stripped and '' in stripped: + before, _, after = stripped.partition('') + reasoning = before.rstrip('\n').rsplit('', 1)[-1].lstrip('\n').strip() + rest = after.lstrip('\n') + message['content'] = f'\n{reasoning}\n\n\n{rest}' + else: + message['content'] = stripped + super()._swift_prepare_inputs(inputs) + + +class Qwen3_8Template(Qwen3_5Template): + reasoning_effort_instructions = { + 'xhigh': ('Reasoning effort is set to xhigh. Please think carefully through the task, validate key ' + 'assumptions, consider plausible alternatives, and prioritize correctness, consistency, ' + 'and clarity in the final answer.'), + 'medium': + '', + 'low': ('Reasoning effort is set to low. Keep your thinking brief and focused, moving directly to the ' + 'conclusion without unnecessary elaboration.'), + } + default_reasoning_effort = 'xhigh' + + def _get_reasoning_instructions(self, inputs: Optional[StdTemplateInputs] = None) -> str: + # HF injects the instruction whenever thinking is on (its jinja default). Follow swift's + # resolved `enable_thinking` so that `--enable_thinking false` (which also emits the + # `\n\n` non-thinking prefix) consistently drops the instruction too. + if not self._get_enable_thinking(inputs): + return '' + reasoning_effort = None + if inputs is not None: + reasoning_effort = inputs.chat_template_kwargs.get('reasoning_effort') + if reasoning_effort is None: + reasoning_effort = self.chat_template_kwargs.get('reasoning_effort') + if reasoning_effort is None: + reasoning_effort = self.default_reasoning_effort + if reasoning_effort not in self.reasoning_effort_instructions: + raise ValueError(f'Unexpected reasoning effort {reasoning_effort}. Supported types are ' + f'{list(self.reasoning_effort_instructions.keys())}.') + return self.reasoning_effort_instructions[reasoning_effort] + + def _get_system(self, inputs: StdTemplateInputs) -> Optional[str]: + system = super()._get_system(inputs) + reasoning_instructions = self._get_reasoning_instructions(inputs) + if not reasoning_instructions: + return system + if system: + return f'{reasoning_instructions}\n\n{system}' + return reasoning_instructions + + def _jinja_encode(self, inputs: StdTemplateInputs): + for message in inputs.messages: + content = message.get('content') + if (message.get('role') != 'assistant' or not isinstance(content, str) + or message.get('reasoning_content') is not None): + continue + if not content.startswith('') or '' not in content: + continue + reasoning_content, _, rest = content.partition('') + message['reasoning_content'] = reasoning_content[len(''):].strip() + message['content'] = rest.lstrip('\n') + return super()._jinja_encode(inputs) + + +register_template( + QwenTemplateMeta( + MLLMTemplateType.qwen3_5, + template_cls=Qwen3_5Template, + default_system=None, + thinking_prefix='\n', + non_thinking_prefix='\n\n\n\n', + agent_template='qwen3_5', + is_thinking=True)) + +register_template( + QwenTemplateMeta( + MLLMTemplateType.qwen3_8, + template_cls=Qwen3_8Template, + default_system=None, + thinking_prefix='\n', + non_thinking_prefix='\n\n\n\n', + agent_template='qwen3_5', + preserve_thinking=True, + is_thinking=True)) + +register_template( + QwenTemplateMeta( + MLLMTemplateType.ovis_ocr2, + template_cls=Qwen3_5Template, + default_system=None, + thinking_prefix='\n', + non_thinking_prefix='\n\n\n\n', + agent_template='qwen3_5', + is_thinking=False)) + + +class Qwen3VLEmbTemplate(Qwen3VLTemplate): + + def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None: + super()._preprocess_inputs(inputs) + if len(inputs.messages) % 2 == 1 and inputs.messages[-1]['role'] != 'assistant': + inputs.messages.append({'role': 'assistant', 'content': ''}) + + +register_template( + QwenTemplateMeta( + MLLMTemplateType.qwen3_vl_emb, + default_system="Represent the user's input.", + suffix=['<|endoftext|>'], + template_cls=Qwen3VLEmbTemplate, + )) + + +class Qwen3VLRerankerTemplate(Qwen3VLTemplate): + instruction = 'Given a search query, retrieve relevant candidates that answer the query.' + + def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None: + super()._preprocess_inputs(inputs) + if inputs.system is not None: + instruction = inputs.system + inputs.system = None + else: + instruction = self.instruction + query = inputs.messages[0]['content'] + document = inputs.messages[1]['content'] + user_message = ': ' + instruction + ':' + query + '\n' + ':' + document + inputs.messages = [{'role': 'user', 'content': user_message}] + return inputs + + +register_template( + QwenTemplateMeta( + MLLMTemplateType.qwen3_vl_reranker, default_system=qwen3_reranker_system, template_cls=Qwen3VLRerankerTemplate)) + + +# ref: trim to hop multiple so WhisperFeatureExtractor matches native HF (floor frames); +# vLLM pad_to_hop_length becomes no-op on pre-trimmed waveforms (GRPO train/rollout align). +def trim_audio_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray: + length = x.shape[-1] + aligned = (length // hop_length) * hop_length + if 0 < aligned < length: + x = x[..., :aligned] + return x + + +class Qwen2_5OmniTemplate(Qwen2_5VLTemplate): + version = 'omni_v2_5' + placeholder_tokens = ['<|IMAGE|>', '<|AUDIO|>', '<|VIDEO|>'] + _requires_mm_token_type_ids = False + + def init_processor(self, processor) -> None: + if processor is None: + return + super().init_processor(processor) + if self.version == 'omni_v2_5': + from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import Qwen2_5OmniProcessorKwargs + default = Qwen2_5OmniProcessorKwargs._defaults + elif self.version == 'omni_v3': + from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import Qwen3OmniMoeProcessorKwargs + default = Qwen3OmniMoeProcessorKwargs._defaults + # Fix: WhisperFeatureExtractor defaults to truncation=True, which silently + # truncates audio longer than 30s. Qwen3 Omni supports variable-length audio, + # so we must disable truncation. See: huggingface/transformers#41473 + default.setdefault('audio_kwargs', {}) + default['audio_kwargs']['truncation'] = False + self.seconds_per_chunk = default['videos_kwargs']['seconds_per_chunk'] + self.position_id_per_seconds = default['videos_kwargs']['position_id_per_seconds'] + self.use_audio_in_video = get_env_args('use_audio_in_video', bool, False) + self.sampling_rate = get_env_args('sampling_rate', int, self.processor.feature_extractor.sampling_rate) + + def _trim_omni_v3_audios(self, audios): + """Trim waveforms to hop-length multiple (omni_v3 only). Matches native HF floor framing.""" + if self.version != 'omni_v3' or not audios: + return audios + hop = self.processor.feature_extractor.hop_length + trimmed = [] + for audio in audios: + if isinstance(audio, tuple): + # train: (wav, 'video'); vllm standalone: (wav, sr) + trimmed.append((trim_audio_to_hop_length(audio[0], hop), audio[1])) + elif isinstance(audio, np.ndarray): + trimmed.append(trim_audio_to_hop_length(audio, hop)) + else: + raise TypeError(f'unexpected audio type {type(audio)!r}; expected ndarray or (ndarray, meta)') + return trimmed + + def _encode_truncated(self, inputs: StdTemplateInputs): + encoded = super()._encode_truncated(inputs) + if self.mode == 'vllm' and inputs.audios: + inputs.audios = self._trim_omni_v3_audios(inputs.audios) + if 'audios' in encoded: + encoded['audios'] = inputs.audios + return encoded + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + from qwen_omni_utils import fetch_image, fetch_video + kwargs = {'image_patch_size': self.processor.image_processor.patch_size} if self.version == 'omni_v3' else {} + sampling_rate = inputs.chat_template_kwargs.get('sampling_rate') + if sampling_rate is None: + sampling_rate = self.sampling_rate + if self.mode == 'vllm': + # https://github.com/modelscope/ms-swift/issues/8445 + inputs.mm_processor_kwargs['do_resize'] = False + if media_type == 'image': + inputs.images[index] = fetch_image({'image': inputs.images[index], **inputs.chat_template_kwargs}, **kwargs) + if self.version == 'omni_v2_5': + return ['<|vision_bos|><|IMAGE|><|vision_eos|>'] + elif self.version == 'omni_v3': + return ['<|vision_start|><|image_pad|><|vision_end|>'] + elif media_type == 'audio': + if self.mode != 'vllm': + inputs.audios[index] = load_audio(inputs.audios[index], sampling_rate) + if self.version == 'omni_v2_5': + return ['<|audio_bos|><|AUDIO|><|audio_eos|>'] + elif self.version == 'omni_v3': + return ['<|audio_start|><|audio_pad|><|audio_end|>'] + elif media_type == 'video': + video = inputs.videos[index] + video_inputs = {'video': video, **inputs.chat_template_kwargs} + if isinstance(video, list): # image list + from qwen_omni_utils import vision_process + video_inputs['sample_fps'] = vision_process.FPS + _video = fetch_video(video_inputs, **kwargs) + if isinstance(_video, torch.Tensor): + _video = _video.to(torch.uint8) + inputs.videos[index] = _video + if self.use_audio_in_video: + if isinstance(video, list): # image list + raise ValueError('image list as video input does not support use_audio_in_video') + audio = load_audio(video, sampling_rate) + if self.mode != 'vllm': + inputs.audios.insert(inputs.audio_idx, (audio, 'video')) + else: + inputs.audios.insert(inputs.audio_idx, audio) + inputs.mm_processor_kwargs['use_audio_in_video'] = True + inputs.audio_idx += 1 + if self.version == 'omni_v2_5': + return ['<|vision_bos|><|audio_bos|><|VIDEO|><|audio_eos|><|vision_eos|>'] + elif self.version == 'omni_v3': + if self.mode == 'vllm': + return ['<|vision_start|><|video_pad|><|vision_end|>'] + else: + return ['<|vision_start|><|audio_start|><|video_pad|><|audio_end|><|vision_end|>'] + if self.version == 'omni_v2_5': + return ['<|vision_bos|><|VIDEO|><|vision_eos|>'] + elif self.version == 'omni_v3': + return ['<|vision_start|><|video_pad|><|vision_end|>'] + + def _get_feat_extract_output_lengths(self, input_lengths): + if self.version == 'omni_v2_5': + return ((input_lengths - 1) // 2 + 1 - 2) // 2 + 1 + elif self.version == 'omni_v3': + input_lengths_leave = input_lengths % 100 + feat_lengths = (input_lengths_leave - 1) // 2 + 1 + return ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 + + def _get_new_tokens_use_audio_in_video(self, i, *, video_grid_thw, video_second_per_grid, audio_lengths, + video_token_id, audio_token_id): + merge_size = self.processor.image_processor.merge_size + grid_thw = video_grid_thw[i] + height = grid_thw[1] // merge_size + width = grid_thw[2] // merge_size + audio_token_indices = torch.arange(audio_lengths[i]) + video_token_indices = torch.arange(grid_thw[0]).reshape(-1, 1, 1) + + video_token_indices = torch.broadcast_to(video_token_indices, + (video_token_indices.shape[0], height, width)).reshape(-1) + video_token_indices = (video_token_indices * video_second_per_grid[i] * self.position_id_per_seconds) + if self.version == 'omni_v2_5': + tokens_per_chunk = int(self.position_id_per_seconds * self.seconds_per_chunk) + video_chunk_indexes = self.processor.get_chunked_index(video_token_indices, tokens_per_chunk) + audio_chunk_indexes = self.processor.get_chunked_index(audio_token_indices, tokens_per_chunk) + + res = [] + for j in range(max(len(video_chunk_indexes), len(audio_chunk_indexes))): + if j < len(video_chunk_indexes): + video_seq_length = video_chunk_indexes[j][1] - video_chunk_indexes[j][0] + res += video_token_id * video_seq_length + if j < len(audio_chunk_indexes): + audio_seq_length = audio_chunk_indexes[j][1] - audio_chunk_indexes[j][0] + res += audio_token_id * audio_seq_length + return res + elif self.version == 'omni_v3': + res = [] + video_data_index, audio_data_index = 0, 0 + while video_data_index < len(video_token_indices) and audio_data_index < len(audio_token_indices): + if video_token_indices[video_data_index] <= audio_token_indices[audio_data_index]: + res += video_token_id + video_data_index += 1 + else: + res += audio_token_id + audio_data_index += 1 + if video_data_index < len(video_token_indices): + res += video_token_id * (len(video_token_indices) - video_data_index) + if audio_data_index < len(audio_token_indices): + res += audio_token_id * (len(audio_token_indices) - audio_data_index) + return res + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = Template._encode(self, inputs) + inputs.audios = self._trim_omni_v3_audios(inputs.audios) + processor = self.processor + video_audios_mask = [] + for i, audio in enumerate(inputs.audios): + if isinstance(audio, tuple) and audio[1] == 'video': + inputs.audios[i] = audio[0] + video_audios_mask.append(True) + else: + video_audios_mask.append(False) + video_audios_mask = torch.tensor(video_audios_mask) + media_inputs = processor( + text='', + audio=inputs.audios or None, + images=inputs.images or None, + videos=inputs.videos or None, + do_resize=False, + return_tensors='pt') + media_inputs.pop('input_ids') + media_inputs.pop('attention_mask') + media_inputs = to_float_dtype(media_inputs, self.model_info.torch_dtype) + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale', None) + config = self.config.thinker_config + # audio + if self.version == 'omni_v3': + audio_token_id = [config.audio_token_id] + else: + audio_token_id = self._tokenize('<|AUDIO|>') + idx_list = findall(input_ids, audio_token_id) + feature_attention_mask = media_inputs.get('feature_attention_mask') + if feature_attention_mask is not None: + audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) + audio_lengths = self._get_feat_extract_output_lengths(audio_feature_lengths) + else: + audio_lengths = None + audio_lengths_origin = audio_lengths + if idx_list: + if self.use_audio_in_video: + audio_lengths = audio_lengths[~video_audios_mask] + + def _get_new_audio_tokens(i): + return audio_token_id * audio_lengths[i] + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_audio_tokens) + + for media_type in ['image', 'video']: + if self.version == 'omni_v3': + token_id = [getattr(config, f'{media_type}_token_id')] + else: + token = f'<|{media_type.upper()}|>' + token_id = self._tokenize(token) + idx_list = findall(input_ids, token_id) + if idx_list: + merge_size = processor.image_processor.merge_size + media_grid_thw = media_inputs.get(f'{media_type}_grid_thw') + if media_type == 'video' and self.use_audio_in_video: + audio_lengths = audio_lengths_origin[video_audios_mask] + video_second_per_grid = media_inputs['video_second_per_grid'] + _get_new_tokens_use_audio_in_video = partial( + self._get_new_tokens_use_audio_in_video, + video_grid_thw=media_grid_thw, + video_second_per_grid=video_second_per_grid, + audio_lengths=audio_lengths, + video_token_id=token_id, + audio_token_id=audio_token_id) + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens_use_audio_in_video) + + else: + + def _get_new_tokens(i): + token_len = (media_grid_thw[i].prod() // (merge_size**2)) + return token_id * token_len + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_tokens) + + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + encoded.update(media_inputs) + return encoded + + def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]: + if not self.is_training: + return inputs + + input_ids = inputs['input_ids'] + input_features = inputs.get('input_features') + feature_attention_mask = inputs.get('feature_attention_mask') + + base_model = self.get_base_model(model) + inputs_embeds = base_model.thinker.model.embed_tokens(input_ids) + thinker_config = model.config.thinker_config + inputs_embeds = self._get_inputs_embeds_hf(inputs_embeds, inputs, model.thinker.visual, self.processor, + thinker_config) + if input_features is None: + if is_deepspeed_enabled() and not is_deepspeed_zero3_enabled(): + # Note: ZeRO-3 still results in hangs; for audio training, please use ZeRO-2. + input_features = input_ids.new_zeros([1, 128, 128], dtype=model.thinker.audio_tower.dtype) + feature_attention_mask = input_ids.new_ones([1, 128], dtype=torch.bool) + audio_res = model.thinker.get_audio_features(input_features, feature_attention_mask) + if hasattr(audio_res, 'last_hidden_state'): + audio_embeds = audio_res.last_hidden_state + else: + audio_embeds = audio_res + inputs_embeds = inputs_embeds + audio_embeds.mean() * 0. + else: + audio_res = model.thinker.get_audio_features(input_features, feature_attention_mask) + if hasattr(audio_res, 'last_hidden_state'): + audio_embeds = audio_res.last_hidden_state + else: + audio_embeds = audio_res + audio_mask = (input_ids == thinker_config.audio_token_index).unsqueeze(-1).expand_as(inputs_embeds) + audio_embeds = audio_embeds.to(inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter(audio_mask, audio_embeds) + + return {'inputs_embeds': inputs_embeds} + + def _get_get_rope_index(self): + return self._get_model().thinker.get_rope_index + + def _get_position_ids(self, inputs: Dict[str, Any]): + if not self.is_training: + return {} + feature_attention_mask = inputs.get('feature_attention_mask') + if feature_attention_mask is not None: + audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) + else: + audio_feature_lengths = None + video_second_per_grid = inputs.pop('video_second_per_grid', None) + input_ids = inputs['input_ids'] + attention_mask = inputs.get('attention_mask_2d') + if attention_mask is None: + attention_mask = inputs.get('attention_mask') + if attention_mask is None: + attention_mask = torch.ones_like(input_ids) + position_ids, _ = self._get_get_rope_index()( + input_ids, + inputs.get('image_grid_thw'), + inputs.get('video_grid_thw'), + attention_mask, + self.use_audio_in_video, + audio_feature_lengths, + video_second_per_grid, + ) + if torch.is_floating_point(position_ids): + position_ids = position_ids.to(torch.int64) + return {'position_ids': self._concat_text_position_ids(position_ids)} + + def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: + res = super()._data_collator_mm_data(batch) + video_second_per_grid = self.gather_list(batch, 'video_second_per_grid') + if video_second_per_grid: + res['video_second_per_grid'] = video_second_per_grid + input_features = [b['input_features'] for b in batch if b.get('input_features') is not None] + feature_attention_mask = [ + b['feature_attention_mask'] for b in batch if b.get('feature_attention_mask') is not None + ] + if input_features: + if self.version == 'omni_v3': + max_length = max(input_feature.shape[-1] for input_feature in input_features) + for i, input_feature in enumerate(input_features): + mask = feature_attention_mask[i] + input_features[i] = F.pad(input_feature, (0, max_length - input_feature.shape[-1])) + feature_attention_mask[i] = F.pad(mask, (0, max_length - mask.shape[-1])) + res['input_features'] = torch.concat(input_features) + res['feature_attention_mask'] = torch.concat(feature_attention_mask) + return res + + def generate(self, model, *args, **kwargs): + if kwargs.get('video_grid_thw') is not None: + kwargs['use_audio_in_video'] = self.use_audio_in_video + return super().generate(model, *args, **kwargs) + + +register_template(QwenTemplateMeta(MLLMTemplateType.qwen2_5_omni, template_cls=Qwen2_5OmniTemplate)) + + +class Qwen3OmniTemplate(Qwen2_5OmniTemplate): + version = 'omni_v3' + norm_bbox = 'norm1000' + placeholder_tokens = ['<|image_pad|>', '<|audio_pad|>', '<|video_pad|>'] + + def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs + + +register_template( + QwenTemplateMeta( + MLLMTemplateType.qwen3_omni, template_cls=Qwen3OmniTemplate, default_system=None, thinking_prefix='\n')) + + +def _qwen3_asr_get_feat_extract_output_lengths(input_lengths): + """Qwen3-ASR Conv2d encoder output length: chunks of 100 frames, 13 tokens each.""" + input_lengths_leave = input_lengths % 100 + feat_lengths = (input_lengths_leave - 1) // 2 + 1 + output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 + return output_lengths + + +class Qwen3ASRTemplate(Template): + placeholder_tokens = ['<|audio_pad|>'] + support_padding_free = True + + def init_env_args(self) -> None: + super().init_env_args() + self.sampling_rate = get_env_args('sampling_rate', int, self.processor.feature_extractor.sampling_rate) + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'audio' + return ['<|audio_start|><|audio_pad|><|audio_end|>'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + sampling_rate = inputs.chat_template_kwargs.get('sampling_rate') + if sampling_rate is None: + sampling_rate = self.sampling_rate + if inputs.audios: + audios = load_batch(inputs.audios, load_func=partial(load_audio, sampling_rate=sampling_rate)) + audio_inputs = self.processor.feature_extractor( + audios, + sampling_rate=sampling_rate, + return_attention_mask=True, + return_tensors='pt', + padding=True, + truncation=False) + audio_inputs['feature_attention_mask'] = audio_inputs.pop('attention_mask') + audio_inputs['input_features'] = to_float_dtype(audio_inputs['input_features'], self.model_info.torch_dtype) + encoded.update(audio_inputs) + + input_ids = encoded['input_ids'] + labels = encoded['labels'] + loss_scale = encoded.get('loss_scale') + audio_token_id = self._tokenize('<|audio_pad|>') + idx_list = findall(input_ids, audio_token_id) + + if idx_list: + feature_attention_mask = audio_inputs.get('feature_attention_mask') + if feature_attention_mask is not None: + audio_feature_lengths = torch.sum(feature_attention_mask, dim=1) + audio_lengths = _qwen3_asr_get_feat_extract_output_lengths(audio_feature_lengths) + + def _get_new_audio_tokens(i): + return audio_token_id * int(audio_lengths[i]) + + input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list, + _get_new_audio_tokens) + encoded['input_ids'] = input_ids + encoded['labels'] = labels + encoded['loss_scale'] = loss_scale + + return encoded + + def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]: + res = super()._data_collator_mm_data(batch) + input_features = [b['input_features'] for b in batch if b.get('input_features') is not None] + feature_attention_mask = [ + b['feature_attention_mask'] for b in batch if b.get('feature_attention_mask') is not None + ] + if input_features: + max_length = max(input_feature.shape[-1] for input_feature in input_features) + for i, input_feature in enumerate(input_features): + mask = feature_attention_mask[i] + input_features[i] = F.pad(input_feature, (0, max_length - input_feature.shape[-1])) + feature_attention_mask[i] = F.pad(mask, (0, max_length - mask.shape[-1])) + res['input_features'] = torch.concat(input_features) + res['feature_attention_mask'] = torch.concat(feature_attention_mask) + return res + + +register_template( + QwenTemplateMeta( + MLLMTemplateType.qwen3_asr, + template_cls=Qwen3ASRTemplate, + # Even without adding a system message, + # the '<|im_start|>system\n<|im_end|>\n' prefix is still present. + # Align with the qwen3_asr template. + system_prefix=None, + default_system=None, + prefix=['<|im_start|>system\n{{SYSTEM}}<|im_end|>\n'])) + + +class Qwen3TTSTemplate(Template): + # ref: https://github.com/QwenLM/Qwen3-TTS/tree/main/finetuning + support_padding_free = False + use_model = True + model_accepts_loss_kwargs = False + + def init_env_args(self) -> None: + super().init_env_args() + self._config_initialized = False + self.target_speaker_embedding = None + # Cache TTS config values for data collation + config = self.config + self._tts_pad_token_id = config.tts_pad_token_id + self._tts_bos_token_id = config.tts_bos_token_id + self._tts_eos_token_id = config.tts_eos_token_id + talker_config = config.talker_config + self._codec_nothink_id = talker_config.codec_nothink_id + self._codec_think_bos_id = talker_config.codec_think_bos_id + self._codec_think_eos_id = talker_config.codec_think_eos_id + self._codec_pad_id = talker_config.codec_pad_id + self._codec_bos_id = talker_config.codec_bos_id + self._codec_eos_token_id = talker_config.codec_eos_token_id + + @staticmethod + def _extract_ref_mel(ref_audio_path: str) -> torch.Tensor: + """Extract mel spectrogram from reference audio for speaker embedding.""" + import librosa + from qwen_tts.core.models.modeling_qwen3_tts import mel_spectrogram + audio, sr = librosa.load(ref_audio_path, sr=None, mono=True) + if audio.ndim > 1: + audio = np.mean(audio, axis=-1) + if sr != 24000: + audio = librosa.resample(audio, orig_sr=sr, target_sr=24000) + mels = mel_spectrogram( + torch.from_numpy(audio.astype(np.float32)).unsqueeze(0), + n_fft=1024, + num_mels=128, + sampling_rate=24000, + hop_size=256, + win_size=1024, + fmin=0, + fmax=12000).transpose(1, 2) # [1, mel_len, 128] + return mels + + def _preprocess_inputs(self, inputs: StdTemplateInputs) -> None: + """Override to skip _add_default_tags since audios here are targets, not inputs.""" + pass + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + # Get text from messages (assistant content) + text = inputs.messages[-1]['content'] if inputs.messages else '' + + # Build TTS text with assistant markers + tts_text = f'<|im_start|>assistant\n{text}' + text_ids = self._tokenize(tts_text) + + # Get audio codes (pre-extracted or online) + audio_codes = inputs.extra_kwargs.get('audio_codes') + if audio_codes is None: + audio_path = inputs.audios[0] if inputs.audios else None + if audio_path: + tts_tokenizer = self.processor.tts_tokenizer + enc_res = tts_tokenizer.encode([audio_path]) + audio_codes = enc_res.audio_codes[0].cpu().tolist() + assert audio_codes is not None, "Either 'audio_codes' or 'audio'/'audios' must be provided in the dataset." + audio_codes = torch.tensor(audio_codes, dtype=torch.long) # [t, 16] + + # Extract mel spectrogram from reference audio + ref_audios = inputs.extra_kwargs.get('ref_audios') + ref_audio_path = ref_audios[0] + assert ref_audio_path is not None, "'ref_audios' must be provided in the dataset." + ref_mel = self._extract_ref_mel(ref_audio_path) # [1, mel_len, 128] + + return { + 'input_ids': text_ids, # dummy for length tracking + 'labels': None, + 'tts_audio_codes': audio_codes, # [codec_len, 16] + 'tts_ref_mel': ref_mel, # [1, mel_len, 128] + } + + def compute_sft_loss(self, model, inputs, num_items_in_batch=None, trainer=None): + """Override to bypass standard label adjustment - TTS loss is computed in forward. + + Combines the talker codec_0 cross-entropy loss with the sub-talker loss + using a fixed weighting factor of 0.3. + """ + # Extract speaker_embedding from ref_mels and cache for checkpoint post-processing + if 'ref_mels' in inputs: + base_model = model.module if hasattr(model, 'module') else model + with torch.no_grad(): + speaker_embedding = base_model.speaker_encoder(inputs['ref_mels'].to(base_model.device).to( + base_model.dtype)).detach() + if self.target_speaker_embedding is None: + self.target_speaker_embedding = speaker_embedding.cpu() + inputs.pop('ref_mels') + inputs['speaker_embedding'] = speaker_embedding + outputs = model(**inputs) + return outputs + + def save_callback(self, model, output_dir): + """Custom save: drop speaker_encoder weights and inject target_speaker_embedding + into codec_embedding.weight[3000].""" + shutil.copytree(model.config.name_or_path, output_dir, dirs_exist_ok=True) + with open(os.path.join(model.config.name_or_path, 'config.json'), 'r', encoding='utf-8') as f: + config_dict = json.load(f) + speaker_name = get_env_args('speaker_name', str, 'speaker_test') + config_dict['tts_model_type'] = 'custom_voice' + config_dict['talker_config']['spk_id'] = {speaker_name: 3000} + config_dict['talker_config']['spk_is_dialect'] = {speaker_name: False} + from safetensors.torch import save_file + from transformers.modeling_utils import unwrap_model + + base_model = unwrap_model(model) + state_dict = {k: v.detach().cpu() for k, v in base_model.state_dict().items()} + + # 1. Drop speaker_encoder keys + keys_to_drop = [k for k in state_dict if k.startswith('speaker_encoder')] + for k in keys_to_drop: + del state_dict[k] + + # 2. Inject target_speaker_embedding into codec_embedding.weight[3000] + emb_key = 'talker.model.codec_embedding.weight' + if self.target_speaker_embedding is not None and emb_key in state_dict: + weight = state_dict[emb_key] + state_dict[emb_key][3000] = self.target_speaker_embedding[0].to(weight.dtype) + + save_file(state_dict, os.path.join(output_dir, 'model.safetensors')) + # Save config + with open(os.path.join(output_dir, 'config.json'), 'w', encoding='utf-8') as f: + json.dump(config_dict, f, indent=2, ensure_ascii=False) + + def data_collator(self, batch: List[Dict[str, Any]], *, padding_to=None) -> Dict[str, Any]: + """Custom TTS data collation - builds dual-channel input format.""" + item_length = [len(b['input_ids']) + b['tts_audio_codes'].shape[0] for b in batch] + max_length = max(item_length) + 8 + b_size, t = len(batch), max_length + + input_ids = torch.zeros((b_size, t, 2), dtype=torch.long) + codec_ids = torch.zeros((b_size, t, 16), dtype=torch.long) + text_embedding_mask = torch.zeros((b_size, t), dtype=torch.bool) + codec_embedding_mask = torch.zeros((b_size, t), dtype=torch.bool) + codec_mask = torch.zeros((b_size, t), dtype=torch.bool) + attention_mask = torch.zeros((b_size, t), dtype=torch.long) + codec_0_labels = torch.full((b_size, t), -100, dtype=torch.long) + + for i, data in enumerate(batch): + text_ids = torch.tensor(data['input_ids']) # [text_len] + audio_codes = data['tts_audio_codes'] # [codec_len, 16] + audio_codec_0 = audio_codes[:, 0] + + text_ids_len = len(text_ids) + codec_ids_len = audio_codec_0.shape[0] + + # === Text channel === + input_ids[i, :3, 0] = text_ids[:3] + input_ids[i, 3:7, 0] = self._tts_pad_token_id + input_ids[i, 7, 0] = self._tts_bos_token_id + input_ids[i, 8:8 + text_ids_len - 3, 0] = text_ids[3:] + input_ids[i, 8 + text_ids_len - 3, 0] = self._tts_eos_token_id + input_ids[i, 8 + text_ids_len - 2:8 + text_ids_len + codec_ids_len, 0] = self._tts_pad_token_id + text_embedding_mask[i, :8 + text_ids_len + codec_ids_len] = True + + # === Codec channel === + input_ids[i, 3:8, 1] = torch.tensor([ + self._codec_nothink_id, + self._codec_think_bos_id, + self._codec_think_eos_id, + 0, # placeholder for speaker embedding + self._codec_pad_id, + ]) + input_ids[i, 8:8 + text_ids_len - 3, 1] = self._codec_pad_id + input_ids[i, 8 + text_ids_len - 3, 1] = self._codec_pad_id + input_ids[i, 8 + text_ids_len - 2, 1] = self._codec_bos_id + input_ids[i, 8 + text_ids_len - 1:8 + text_ids_len - 1 + codec_ids_len, 1] = audio_codec_0 + input_ids[i, 8 + text_ids_len - 1 + codec_ids_len, 1] = self._codec_eos_token_id + + # === Labels (codec layer 0) === + codec_0_labels[i, 8 + text_ids_len - 1:8 + text_ids_len - 1 + codec_ids_len] = audio_codec_0 + codec_0_labels[i, 8 + text_ids_len - 1 + codec_ids_len] = self._codec_eos_token_id + + # === Sub-talker codec IDs === + codec_ids[i, 8 + text_ids_len - 1:8 + text_ids_len - 1 + codec_ids_len, :] = audio_codes + + # === Masks === + codec_embedding_mask[i, 3:8 + text_ids_len + codec_ids_len] = True + codec_embedding_mask[i, 6] = False # speaker embedding position + codec_mask[i, 8 + text_ids_len - 1:8 + text_ids_len - 1 + codec_ids_len] = True + attention_mask[i, :8 + text_ids_len + codec_ids_len] = True + + ref_mels = torch.cat([data['tts_ref_mel'] for data in batch], dim=0) + + return { + 'input_ids': input_ids, + 'ref_mels': ref_mels, + 'attention_mask': attention_mask, + 'text_embedding_mask': text_embedding_mask.unsqueeze(-1), + 'codec_embedding_mask': codec_embedding_mask.unsqueeze(-1), + 'labels': codec_0_labels, + 'codec_0_labels': codec_0_labels, + 'codec_ids': codec_ids, + 'codec_mask': codec_mask, + } + + +register_template(QwenTemplateMeta( + MLLMTemplateType.qwen3_tts, + template_cls=Qwen3TTSTemplate, + default_system=None, +)) + + +class Ovis1_6Template(Template): + skip_prompt = False + use_model = True + + def init_env_args(self): + super().init_env_args() + self.max_partition = get_env_args('max_partition', int, 9) + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + assert media_type == 'image' + return [[-200], '\n'] + + def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]: + encoded = super()._encode(inputs) + images = inputs.images + input_ids = encoded['input_ids'] + labels = encoded['labels'] + idx_list = findall(input_ids, [-200]) + added_tokens_len = 0 + pixel_values = [] + for i, idx in enumerate(idx_list): + raw_pixel_values, image_placeholders = self.model.visual_tokenizer.preprocess_image( + images[i], max_partition=self.max_partition) + input_ids = input_ids[:idx] + image_placeholders + input_ids[idx + 1:] + if labels is not None: + labels = labels[:idx] + [-100] * len(image_placeholders) + labels[idx + 1:] + pixel_values.append(raw_pixel_values) + added_tokens_len += len(image_placeholders) - 1 + dtype = self.model.visual_tokenizer.dtype + if pixel_values: + pixel_values = torch.cat(pixel_values, dim=0).to(dtype) + else: + pixel_values = torch.zeros((1, 3, 384, 384), dtype=dtype) # dummpy + encoded.update({'input_ids': input_ids, 'labels': labels}) + encoded['pixel_values'] = [pixel_values] + return encoded + + def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]: + padding_side = self.padding_side if self.is_training else 'left' + if self.max_length is not None: + model.config.multimodal_max_length = self.max_length + input_ids = inputs['input_ids'] + labels = inputs.get('labels') + if labels is None: + labels = input_ids.new_full(input_ids.shape, -100) + _, inputs_embeds, labels, attention_mask = model.merge_multimodal( + text_input_ids=input_ids, + text_attention_masks=torch.ones_like(input_ids), # not use, only compat + text_labels=labels, + pixel_values=inputs['pixel_values'], + left_padding=padding_side == 'left') + if inputs.get('labels') is None: + labels = None + return {'inputs_embeds': inputs_embeds, 'labels': labels, 'attention_mask': attention_mask} + + def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]: + pixel_values = self.gather_list(batch, 'pixel_values') + res = super()._data_collator(batch, padding_to=padding_to) + res['pixel_values'] = pixel_values + return res + + +register_template( + TemplateMeta( + MLLMTemplateType.ovis1_6, + prefix=[''], + prompt=['user\n{{QUERY}}\nmodel\n'], + chat_sep=['\n'], + suffix=[''], + system_prefix=['system\n{{SYSTEM}}\n'], + template_cls=Ovis1_6Template, + )) + +register_template( + Llama3TemplateMeta( + MLLMTemplateType.ovis1_6_llama3, + default_system='You are a helpful and honest multimodal assistant.', + template_cls=Ovis1_6Template, + agent_template=None, + )) + + +class Ovis2Template(Ovis1_6Template): + placeholder_tokens = ['<|image_pad|>', '<|video_pad|>'] + NFRAMES = 12 + + def init_env_args(self): + super().init_env_args() + self.nframes = get_env_args('nframes', int, self.NFRAMES) + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + if self.mode == 'vllm': + return ['\n'] + return [[-200], '\n'] + elif media_type == 'video': + inputs.images = load_video_ovis2(inputs.videos[index], self.nframes) + return [[-200] * self.nframes, '\n'] + + +register_template(QwenTemplateMeta( + MLLMTemplateType.ovis2, + template_cls=Ovis2Template, +)) + + +class Ovis2_5Template(Template): + use_model = True + skip_prompt = False + support_padding_free = True + + def init_env_args(self) -> None: + super().init_env_args() + self.min_pixels = get_env_args('min_pixels', int, 448 * 448) + self.max_pixels = get_env_args('max_pixels', int, 1344 * 1792) + self.video_max_pixels = get_env_args('video_max_pixels', int, 896 * 896) + self.num_frames = get_env_args('num_frames', int, 8) + + def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int, + inputs: StdTemplateInputs) -> List[Context]: + if media_type == 'image': + if self.mode == 'vllm': + return [''] + else: + return [[-200], '\n'] + elif media_type == 'video': + if self.mode == 'vllm': + return ['