eagle3 vlm训练，数据预处理过程很慢，而且内存占用很高

[wincle]

[yghstill]

fix in #176

[wincle]

更新后问题似乎并没有修复，我用了1000条数据测试，8卡h20训练，模型是qwen-vl-2b，max-length 8192，数据加载有些慢，而且会出异常，我感觉是内存超了，我把max-length缩小到 1024是能train 起来的，下面是异常的日志。

这是我的训练脚本
export CONFIG_DIR=angelslim/compressor/speculative/train/configs
export TARGET_MODEL_NAME_OR_PATH=ocr_model/2b_model_ocr/
export DRAFT_MODEL_CONFIG_PATH=$CONFIG_DIR/qwen3-vl-2b-eagle3-mrope.json
export TRAIN_DATA_PATH=/UserData/dataset/ocr/nature/eval_eagle3.jsonl
export EVAL_DATA_PATH=/UserData/dataset/ocr/nature/eval_eagle3.jsonl
export OUTPUT_DIR=ocr_model/2b_model_ocr_eagle3
export MODEL_MAX_LENGTH=1024
export CHAT_TEMPLATE_TYPE=qwen3_vl
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

torchrun --nproc_per_node=8 tools/train_eagle3_online.py
--modal_type VLM
--target_model_name_or_path $TARGET_MODEL_NAME_OR_PATH
--draft_model_config_path $DRAFT_MODEL_CONFIG_PATH
--train_data_path $TRAIN_DATA_PATH
--eval_data_path $EVAL_DATA_PATH
--output_dir $OUTPUT_DIR
--num_train_epochs 20
--per_device_train_batch_size 1
--per_device_eval_batch_size 1
--gradient_accumulation_steps 1
--save_strategy "steps"
--save_steps 4000
--learning_rate 1e-4
--weight_decay 0.0
--warmup_ratio 0.1
--lr_scheduler_type "constant"
--logging_steps 20
--model_max_length $MODEL_MAX_LENGTH
--deepspeed $CONFIG_DIR/deepspeed_zero3.json
--chat_template_type $CHAT_TEMPLATE_TYPE
--report_to none
--run_name qwen3-2b-eagle3-angelslim \

异常日志：
$./train_vlm_eagle.sh
W1216 10:23:55.664000 6390 site-packages/torch/distributed/run.py:774]
W1216 10:23:55.664000 6390 site-packages/torch/distributed/run.py:774] *****************************************
W1216 10:23:55.664000 6390 site-packages/torch/distributed/run.py:774] Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed.
W1216 10:23:55.664000 6390 site-packages/torch/distributed/run.py:774] *****************************************
[Single Process] Loading model to device: cuda:7
Loading target model with hf backend...
[Single Process] Loading model to device: cuda:0
[Single Process] Loading model to device: cuda:1
[Single Process] Loading model to device: cuda:4
[Single Process] Loading model to device: cuda:3
[Single Process] Loading model to device: cuda:5
[Single Process] Loading model to device: cuda:2
[Single Process] Loading model to device: cuda:6
The module name (originally ) is not a valid Python identifier. Please rename the original module to avoid import issues.
The module name (originally ) is not a valid Python identifier. Please rename the original module to avoid import issues.
The module name (originally ) is not a valid Python identifier. Please rename the original module to avoid import issues.
The module name (originally ) is not a valid Python identifier. Please rename the original module to avoid import issues.
The module name (originally ) is not a valid Python identifier. Please rename the original module to avoid import issues.
The module name (originally ) is not a valid Python identifier. Please rename the original module to avoid import issues.
The module name (originally ) is not a valid Python identifier. Please rename the original module to avoid import issues.
The module name (originally ) is not a valid Python identifier. Please rename the original module to avoid import issues.
Unrecognized keys in rope_scaling for 'rope_type'='default': {'mrope_section', 'mrope_interleaved'}
Target model loaded successfully
Loading draft model...
Unrecognized keys in rope_scaling for 'rope_type'='default': {'mrope_section', 'mrope_interleaved'}
draft_model_config: LlamaConfig {
"architectures": [
"Eagle3LlamaForCausalLM"
],
"attention_bias": false,
"attention_dropout": 0.0,
"bos_token_id": 151643,
"draft_vocab_size": 32000,
"dtype": "bfloat16",
"eos_token_id": 151645,
"head_dim": 128,
"hidden_act": "silu",
"hidden_size": 2048,
"image_token_id": 151655,
"initializer_range": 0.02,
"intermediate_size": 9728,
"max_position_embeddings": 262144,
"mlp_bias": false,
"modal_type": "VLM",
"model_type": "llama",
"num_attention_heads": 32,
"num_hidden_layers": 1,
"num_key_value_heads": 8,
"pretraining_tp": 1,
"rms_norm_eps": 1e-06,
"rope_scaling": {
"mrope_interleaved": true,
"mrope_section": [
24,
20,
20
],
"rope_type": "default",
"type": "default"
},
"rope_theta": 5000000,
"tie_word_embeddings": true,
"transformers_version": "4.57.1",
"use_cache": true,
"video_token_id": 151656,
"vision_end_token_id": 151653,
"vision_start_token_id": 151652,
"vocab_size": 151936
}

Unrecognized keys in rope_scaling for 'rope_type'='default': {'mrope_section', 'mrope_interleaved'}
Unrecognized keys in rope_scaling for 'rope_type'='default': {'mrope_interleaved', 'mrope_section'}
Unrecognized keys in rope_scaling for 'rope_type'='default': {'mrope_interleaved', 'mrope_section'}
Unrecognized keys in rope_scaling for 'rope_type'='default': {'mrope_interleaved', 'mrope_section'}
Unrecognized keys in rope_scaling for 'rope_type'='default': {'mrope_interleaved', 'mrope_section'}
Unrecognized keys in rope_scaling for 'rope_type'='default': {'mrope_interleaved', 'mrope_section'}
torch_dtype is deprecated! Use dtype instead!
torch_dtype is deprecated! Use dtype instead!
torch_dtype is deprecated! Use dtype instead!
torch_dtype is deprecated! Use dtype instead!
torch_dtype is deprecated! Use dtype instead!
torch_dtype is deprecated! Use dtype instead!
torch_dtype is deprecated! Use dtype instead!
torch_dtype is deprecated! Use dtype instead!
Draft model loaded successfully
Creating training and evaluation datasets with chat template type: qwen3_vl...
Processing conversations (num_proc=16): 0%| | 0/947 [00:00<?, ? examples/s]Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8160]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8742]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8800]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[9312]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8769]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[10829]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8769]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8769]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8260]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8160]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8775]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8580]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[12288]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8160]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8532]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8265]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Error processing conversation: Mismatch in image token count between text and input_ids. Got ids=[8083] and text=[8769]. Likely due to truncation='max_length'. Please disable truncation or increase max_length.
Processing conversations (num_proc=16): 0%| | 0/947 [01:25<?, ? examples/s]
Traceback (most recent call last):
File "/usr/lib/python3.12/site-packages/angelslim-0.0.0.dev0-py3.12.egg/angelslim/compressor/speculative/train/data/dataset_builder/online_vlm_dataset_builder.py", line 91, in build_dataset
processed_ds = ds.map(
^^^^^^^
File "/usr/lib/python3.12/site-packages/datasets/arrow_dataset.py", line 557, in wrapper
out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/site-packages/datasets/arrow_dataset.py", line 3171, in map
for rank, done, content in iflatmap_unordered(
^^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/site-packages/datasets/utils/py_utils.py", line 721, in iflatmap_unordered
raise RuntimeError(
RuntimeError: One of the subprocesses has abruptly died during map operation.To debug the error, disable multiprocessing.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/ProjectRoot/AngelSlim/tools/train_eagle3_online.py", line 413, in
train()
File "/ProjectRoot/AngelSlim/tools/train_eagle3_online.py", line 314, in train
dataset_manager.create_online_datasets()
File "/usr/lib/python3.12/site-packages/angelslim-0.0.0.dev0-py3.12.egg/angelslim/compressor/speculative/train/data/dataset.py", line 154, in create_online_datasets
return self._create_online_datasets()
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/site-packages/angelslim-0.0.0.dev0-py3.12.egg/angelslim/compressor/speculative/train/data/dataset.py", line 174, in _create_online_datasets
train_dataset = self.online_dataset_builder.build_dataset(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/site-packages/angelslim-0.0.0.dev0-py3.12.egg/angelslim/compressor/speculative/train/data/dataset_builder/online_vlm_dataset_builder.py", line 112, in build_dataset
raise RuntimeError(f"Dataset building failed for {datapath}") from e
RuntimeError: Dataset building failed for ['/UserData/dataset/ocr/nature/eval_eagle3.jsonl']
Processing conversations (num_proc=16): 0%| | 0/947 [01:53<?, ? examples/s]
W1216 10:26:12.854000 6390 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 6467 closing signal SIGTERM
W1216 10:26:12.854000 6390 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 6468 closing signal SIGTERM
W1216 10:26:12.855000 6390 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 6470 closing signal SIGTERM
W1216 10:26:12.856000 6390 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 6471 closing signal SIGTERM
W1216 10:26:12.857000 6390 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 6472 closing signal SIGTERM
W1216 10:26:12.857000 6390 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 6473 closing signal SIGTERM
W1216 10:26:12.858000 6390 site-packages/torch/distributed/elastic/multiprocessing/api.py:900] Sending process 6474 closing signal SIGTERM
E1216 10:26:13.424000 6390 site-packages/torch/distributed/elastic/multiprocessing/api.py:874] failed (exitcode: 1) local_rank: 2 (pid: 6469) of binary: /usr/bin/python3
Traceback (most recent call last):
File "/usr/bin/torchrun", line 33, in
sys.exit(load_entry_point('torch==2.8.0', 'console_scripts', 'torchrun')())
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/init.py", line 357, in wrapper
return f(*args, **kwargs)
^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/site-packages/torch/distributed/run.py", line 901, in main
run(args)
File "/usr/lib/python3.12/site-packages/torch/distributed/run.py", line 892, in run
elastic_launch(
File "/usr/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 143, in call
return launch_agent(self._config, self._entrypoint, list(args))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib/python3.12/site-packages/torch/distributed/launcher/api.py", line 277, in launch_agent
raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError:

tools/train_eagle3_online.py FAILED

Failures:
<NO_OTHER_FAILURES>

Root Cause (first observed failure):
[0]:
time : 2025-12-16_10:26:12
host : dl2-prod-instance-54-job-0-398cd8-0
rank : 2 (local_rank: 2)
exitcode : 1 (pid: 6469)
error_file: <N/A>
traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html

[yghstill]

@wincle 好的，我们再定位下问题