Skip to content

Commit 518db42

Browse files
Schnabel-8PanAndy
Schnabel-8
authored and
PanAndy
committed
(feat): add support for multi-modal distill.
1 parent a73e846 commit 518db42

9 files changed

Lines changed: 500 additions & 4 deletions

File tree

‎examples/qwen2.5-1.5B-distill_ds/distill_zero3.yaml‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,6 @@ checkpoint_config:
1818
type: file_system
1919
output_dir: /data/oss_bucket_0/chuye/roll/distill/models/zero3
2020

21-
track_with: ml_tracker
22-
2321
save_steps: 100
2422
logging_steps: 1
2523
resume_from_checkpoint: false

‎examples/qwen2.5-7B-distill_megatron/distill_megatron.yaml‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,6 @@ checkpoint_config:
1212
type: file_system
1313
output_dir: /data/oss_bucket_0/chuye/roll/distill/models/megatron_dp_pp
1414

15-
track_with: ml_tracker
16-
1715
save_steps: 100
1816
logging_steps: 1
1917
resume_from_checkpoint: false

‎examples/qwen2.5-vl-7B-distill/distill_vl_megatron.yaml‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
hydra:
2+
run:
3+
dir: .
4+
output_subdir: null
5+
6+
exp_name: "distill_vl_megatron"
7+
seed: 42
8+
logging_dir: ./output/logs
9+
output_dir: ./output
10+
11+
checkpoint_config:
12+
type: file_system
13+
output_dir: /data/cpfs_0/rl_examples/models/${exp_name}
14+
15+
save_steps: 100
16+
logging_steps: 1
17+
resume_from_checkpoint: false
18+
19+
student_pretrain: Qwen/Qwen2.5-VL-7B-Instruct
20+
teacher_pretrain: Qwen/Qwen2.5-VL-32B-Instruct
21+
22+
# distill config
23+
distill_loss_weight: 0.85
24+
kd_objective: forward_kl
25+
distill_on_prompt: True
26+
27+
sequence_length: 1024
28+
max_grad_norm: 1.0
29+
30+
student:
31+
model_args:
32+
attn_implementation: fa2
33+
disable_gradient_checkpointing: false
34+
dtype: bf16
35+
model_type: ~
36+
training_args:
37+
learning_rate: 2.0e-5
38+
lr_scheduler_type: constant
39+
per_device_train_batch_size: 4
40+
gradient_accumulation_steps: 1
41+
warmup_steps: 0
42+
num_train_epochs: 1
43+
max_steps: 1000
44+
data_args:
45+
template: qwen2-vl
46+
# use leonardPKU/GEOQA_R1V_Train_8K as dataset
47+
# download to ./data/geoqa_data from https://huggingface.co/datasets/leonardPKU/GEOQA_R1V_Train_8K
48+
file_name: data/geoqa_data/
49+
dataset_dir: ./
50+
preprocessing_num_workers: 16
51+
strategy_args:
52+
strategy_name: megatron_train
53+
strategy_config:
54+
tensor_model_parallel_size: 2
55+
pipeline_model_parallel_size: 2
56+
use_distributed_optimizer: true
57+
recompute_granularity: full
58+
device_mapping: list(range(0,8))
59+
60+
teacher:
61+
model_args:
62+
attn_implementation: fa2
63+
disable_gradient_checkpointing: true
64+
dtype: bf16
65+
data_args:
66+
template: qwen2-vl
67+
strategy_args:
68+
strategy_name: megatron_infer
69+
strategy_config:
70+
tensor_model_parallel_size: 2
71+
pipeline_model_parallel_size: 2
72+
bf16: true
73+
device_mapping: list(range(0,8))
74+
75+
system_envs:
76+
RAY_PROFILING: "0"

‎examples/qwen2.5-vl-7B-distill/distill_vl_zero3.yaml‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
defaults:
2+
- ../config/deepspeed_zero@_here_
3+
- ../config/deepspeed_zero2@_here_
4+
- ../config/deepspeed_zero3@_here_
5+
- ../config/deepspeed_zero3_cpuoffload@_here_
6+
7+
hydra:
8+
run:
9+
dir: .
10+
output_subdir: null
11+
12+
exp_name: "distill_vl_zero3"
13+
seed: 42
14+
logging_dir: ./output/logs
15+
output_dir: ./output
16+
17+
checkpoint_config:
18+
type: file_system
19+
output_dir: /data/cpfs_0/rl_examples/models/${exp_name}
20+
21+
save_steps: 100
22+
logging_steps: 1
23+
resume_from_checkpoint: false
24+
25+
student_pretrain: Qwen/Qwen2.5-VL-7B-Instruct
26+
teacher_pretrain: Qwen/Qwen2.5-VL-32B-Instruct
27+
28+
# distill config
29+
distill_loss_weight: 0.85
30+
kd_objective: forward_kl
31+
distill_on_prompt: True
32+
33+
sequence_length: 1024
34+
max_grad_norm: 1.0
35+
36+
student:
37+
model_args:
38+
attn_implementation: fa2
39+
disable_gradient_checkpointing: false
40+
dtype: bf16
41+
model_type: ~
42+
training_args:
43+
learning_rate: 2.0e-5
44+
weight_decay: 1.0e-2
45+
lr_scheduler_type: constant
46+
per_device_train_batch_size: 1
47+
gradient_accumulation_steps: 1
48+
warmup_steps: 0
49+
num_train_epochs: 1
50+
data_args:
51+
template: qwen2-vl
52+
# use leonardPKU/GEOQA_R1V_Train_8K as dataset
53+
# download to ./data/geoqa_data from https://huggingface.co/datasets/leonardPKU/GEOQA_R1V_Train_8K
54+
file_name: data/geoqa_data/
55+
dataset_dir: ./
56+
preprocessing_num_workers: 16
57+
58+
strategy_args:
59+
strategy_name: deepspeed_train
60+
strategy_config: ${deepspeed_zero3}
61+
device_mapping: list(range(0,8))
62+
63+
teacher:
64+
model_args:
65+
attn_implementation: fa2
66+
disable_gradient_checkpointing: true
67+
dtype: bf16
68+
data_args:
69+
template: qwen2-vl
70+
strategy_args:
71+
strategy_name: deepspeed_infer
72+
strategy_config: ${deepspeed_zero3}
73+
device_mapping: list(range(0,8))
74+
75+
system_envs:
76+
RAY_PROFILING: "0"

‎examples/qwen2.5-vl-7B-distill/run_distill_vl_ds_pipeline.sh‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
2+
#!/bin/bash
3+
set +x
4+
5+
CONFIG_PATH=$(basename $(dirname $0))
6+
python examples/start_distill_vl_pipeline.py --config_path $CONFIG_PATH --config_name distill_vl_zero3

‎examples/qwen2.5-vl-7B-distill/run_distill_vl_megatron_pipeline.sh‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
2+
#!/bin/bash
3+
set +x
4+
5+
CONFIG_PATH=$(basename $(dirname $0))
6+
python examples/start_distill_vl_pipeline.py --config_path $CONFIG_PATH --config_name distill_vl_megatron

‎examples/start_distill_vl_pipeline.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import argparse
2+
3+
from dacite import from_dict, Config
4+
from hydra import compose, initialize
5+
from omegaconf import OmegaConf
6+
7+
from roll.distributed.scheduler.initialize import init
8+
from roll.pipeline.distill.distill_config import DistillConfig
9+
10+
from roll.pipeline.distill.distill_vlm_pipeline import DistillVLMPipeline
11+
12+
13+
def main():
14+
parser = argparse.ArgumentParser()
15+
parser.add_argument("--config_path", help="The path of the main configuration file", default="config")
16+
parser.add_argument(
17+
"--config_name", help="The name of the main configuration file (without extension).", default="sppo_config"
18+
)
19+
args = parser.parse_args()
20+
21+
initialize(config_path=args.config_path, job_name="app")
22+
cfg = compose(config_name=args.config_name)
23+
24+
print(OmegaConf.to_yaml(cfg, resolve=True))
25+
26+
distill_config = from_dict(data_class=DistillConfig, data=OmegaConf.to_container(cfg, resolve=True))
27+
28+
init()
29+
30+
pipeline = DistillVLMPipeline(pipeline_config=distill_config)
31+
32+
pipeline.run()
33+
34+
35+
if __name__ == "__main__":
36+
main()

‎roll/datasets/collator.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -218,3 +218,12 @@ def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
218218
batch[key] = np.empty(len(batch[key]), dtype=object)
219219
batch[key][:] = batch[key]
220220
return batch
221+
222+
@dataclass
223+
class DataCollatorWithPaddingForMMWithLabels(DataCollatorWithPaddingForMM):
224+
def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
225+
batch = super().__call__(features)
226+
labels = batch["input_ids"].clone()
227+
labels[batch["attention_mask"] == 0] = -100
228+
batch["labels"] = labels
229+
return batch

0 commit comments

Comments
 (0)