mirror of
https://github.com/volcengine/verl.git
synced 2025-10-20 13:43:50 +08:00
[trainer] feat: VL support freeze vision model (#3178)
### What does this PR do? vl model support freeze vision model issue: [2526](https://github.com/volcengine/verl/issues/2526) > Add **concise** overview of what this PR aims to achieve or accomplish. Reference related GitHub issues and PRs that help with the review. ### Checklist Before Starting - [ ] Search for similar PRs. Paste at least one query link here: ... - [ ] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI) - `{modules}` include `fsdp`, `megatron`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data` - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]` - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test` - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title. - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching` ### Test > For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc. qwen2_vl_7b_function_rm_1756093906 is vision freeze mode <img width="4374" height="2086" alt="image" src="https://github.com/user-attachments/assets/107772e4-039d-4ec5-b193-54688f4a7176" /> ### API and Usage Example > Demonstrate how the API changes if any, and provide usage example(s) if possible. ```python # Add code snippet or script demonstrating how to use this ``` ### Design & Code Changes > Demonstrate the high-level design if this PR is complex, and list the specific changes. ### Checklist Before Submitting > [!IMPORTANT] > Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review. - [ ] Read the [Contribute Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md). - [ ] Apply [pre-commit checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always` - [ ] Add / Update [the documentation](https://github.com/volcengine/verl/tree/main/docs). - [ ] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/volcengine/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ... - [ ] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ). (If not accessible, please try [the Feishu group (飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).) --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Mighten Dai <mighten@outlook.com>
This commit is contained in:
maijia-cwh
GitHub
47
examples/grpo_trainer/run_qwen2_5_vl-7b_freeze_vision.sh
Normal file
47
examples/grpo_trainer/run_qwen2_5_vl-7b_freeze_vision.sh
Normal file
@ -0,0 +1,47 @@
|
||||
set -x
|
||||
ENGINE=${1:-vllm}
|
||||
|
||||
python3 -m verl.trainer.main_ppo \
|
||||
algorithm.adv_estimator=grpo \
|
||||
data.train_files=$HOME/data/geo3k/train.parquet \
|
||||
data.val_files=$HOME/data/geo3k/test.parquet \
|
||||
data.train_batch_size=512 \
|
||||
data.max_prompt_length=1024 \
|
||||
data.max_response_length=2048 \
|
||||
data.filter_overlong_prompts=True \
|
||||
data.truncation='error' \
|
||||
data.image_key=images \
|
||||
actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
actor_rollout_ref.actor.optim.lr=1e-6 \
|
||||
actor_rollout_ref.actor.freeze_vision_tower=True \
|
||||
actor_rollout_ref.model.use_remove_padding=True \
|
||||
actor_rollout_ref.actor.ppo_mini_batch_size=128 \
|
||||
actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \
|
||||
actor_rollout_ref.actor.use_kl_loss=True \
|
||||
actor_rollout_ref.actor.kl_loss_coef=0.01 \
|
||||
actor_rollout_ref.actor.kl_loss_type=low_var_kl \
|
||||
actor_rollout_ref.actor.entropy_coeff=0 \
|
||||
actor_rollout_ref.model.enable_gradient_checkpointing=True \
|
||||
actor_rollout_ref.actor.fsdp_config.param_offload=False \
|
||||
actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \
|
||||
actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 \
|
||||
actor_rollout_ref.rollout.tensor_model_parallel_size=2 \
|
||||
actor_rollout_ref.rollout.name=$ENGINE \
|
||||
+actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \
|
||||
actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \
|
||||
actor_rollout_ref.rollout.enable_chunked_prefill=False \
|
||||
actor_rollout_ref.rollout.enforce_eager=False \
|
||||
actor_rollout_ref.rollout.free_cache_engine=True \
|
||||
actor_rollout_ref.rollout.n=5 \
|
||||
actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=20 \
|
||||
actor_rollout_ref.ref.fsdp_config.param_offload=True \
|
||||
algorithm.use_kl_in_reward=False \
|
||||
trainer.critic_warmup=0 \
|
||||
trainer.logger='["console","wandb"]' \
|
||||
trainer.project_name='verl_grpo_example_geo3k' \
|
||||
trainer.experiment_name='qwen2_5_vl_7b_function_rm' \
|
||||
trainer.n_gpus_per_node=8 \
|
||||
trainer.nnodes=1 \
|
||||
trainer.save_freq=20 \
|
||||
trainer.test_freq=5 \
|
||||
trainer.total_epochs=15 $@
|
||||
@ -62,6 +62,7 @@ actor_rollout_ref:
|
||||
clip_ratio: 0.2
|
||||
clip_ratio_low: 0.2
|
||||
clip_ratio_high: 0.2
|
||||
freeze_vision_tower: false
|
||||
policy_loss:
|
||||
_target_: verl.workers.config.PolicyLossConfig
|
||||
loss_mode: vanilla
|
||||
|
||||
@ -47,6 +47,7 @@ actor_rollout_ref:
|
||||
clip_ratio: 0.2
|
||||
clip_ratio_low: 0.2
|
||||
clip_ratio_high: 0.2
|
||||
freeze_vision_tower: false
|
||||
policy_loss:
|
||||
_target_: verl.workers.config.PolicyLossConfig
|
||||
loss_mode: vanilla
|
||||
|
||||
@ -38,6 +38,9 @@ clip_ratio_low: 0.2
|
||||
# Upper bound for asymmetric clipping (used in dual-clip PPO)
|
||||
clip_ratio_high: 0.2
|
||||
|
||||
# Whether to freeze vision model, if set true, it will be freeze vision model
|
||||
freeze_vision_tower: false
|
||||
|
||||
# policy loss config
|
||||
policy_loss:
|
||||
|
||||
|
||||
@ -101,6 +101,7 @@ class ActorConfig(BaseConfig):
|
||||
clip_ratio: float = 0.2
|
||||
clip_ratio_low: float = 0.2
|
||||
clip_ratio_high: float = 0.2
|
||||
freeze_vision_tower: bool = False
|
||||
policy_loss: PolicyLossConfig = field(default_factory=PolicyLossConfig)
|
||||
clip_ratio_c: float = 3.0
|
||||
loss_agg_mode: str = "token-mean"
|
||||
|
||||
@ -119,6 +119,19 @@ def get_sharding_strategy(device_mesh):
|
||||
return sharding_strategy
|
||||
|
||||
|
||||
def get_vl_model_vision_tower(vl_model_instance):
|
||||
"""
|
||||
Util to extract Vision Tower from a VL model instance
|
||||
"""
|
||||
if hasattr(vl_model_instance, "model") and hasattr(vl_model_instance.model, "visual"):
|
||||
# transformers >= 4.52.0
|
||||
return vl_model_instance.model.visual
|
||||
elif hasattr(vl_model_instance, "visual"):
|
||||
# transformers < 4.52.0
|
||||
return vl_model_instance.visual
|
||||
return None
|
||||
|
||||
|
||||
class ActorRolloutRefWorker(Worker, DistProfilerExtension):
|
||||
"""
|
||||
This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy
|
||||
@ -175,6 +188,7 @@ class ActorRolloutRefWorker(Worker, DistProfilerExtension):
|
||||
self._is_actor = self.role in ["actor", "actor_rollout", "actor_rollout_ref"]
|
||||
self._is_rollout = self.role in ["rollout", "actor_rollout", "actor_rollout_ref"]
|
||||
self._is_ref = self.role in ["ref", "actor_rollout_ref"]
|
||||
self.use_orig_params = self.config.actor.fsdp_config.get("use_orig_params", False)
|
||||
|
||||
# TODO(haibin.lin):
|
||||
# As of now the type of config is DictConfig, if we assign config.profiler with ProfilerConfig,
|
||||
@ -395,6 +409,19 @@ class ActorRolloutRefWorker(Worker, DistProfilerExtension):
|
||||
"bias": "none",
|
||||
}
|
||||
actor_module = get_peft_model(actor_module, LoraConfig(**lora_config))
|
||||
|
||||
self.use_orig_params = fsdp_config.get("use_orig_params", False)
|
||||
if self.config.actor.get("freeze_vision_tower", False):
|
||||
vision_tower = get_vl_model_vision_tower(actor_module)
|
||||
if vision_tower is not None:
|
||||
vision_tower.requires_grad_(False)
|
||||
self.use_orig_params = True
|
||||
if self.rank == 0:
|
||||
print("[actor model] Vision tower is set to not trainable.")
|
||||
else:
|
||||
if self.rank == 0:
|
||||
print("[actor model] No vision tower found.")
|
||||
|
||||
torch.distributed.barrier()
|
||||
|
||||
if self.rank == 0:
|
||||
@ -447,7 +474,7 @@ class ActorRolloutRefWorker(Worker, DistProfilerExtension):
|
||||
mixed_precision=mixed_precision,
|
||||
sync_module_states=True,
|
||||
device_mesh=self.device_mesh,
|
||||
use_orig_params=fsdp_config.get("use_orig_params", False),
|
||||
use_orig_params=self.use_orig_params,
|
||||
forward_prefetch=fsdp_config.get("forward_prefetch", False),
|
||||
)
|
||||
elif fsdp_strategy == "fsdp2":
|
||||
@ -1121,6 +1148,7 @@ class CriticWorker(Worker, DistProfilerExtension):
|
||||
f"ppo_micro_batch_size_per_gpu {self.config.ppo_micro_batch_size_per_gpu}"
|
||||
)
|
||||
self._is_lora = self.config.model.get("lora_rank", 0) > 0
|
||||
self.use_orig_params = self.config.model.fsdp_config.get("use_orig_params", False)
|
||||
|
||||
def _build_critic_model_optimizer(self, config):
|
||||
# the following line is necessary
|
||||
@ -1248,12 +1276,24 @@ class CriticWorker(Worker, DistProfilerExtension):
|
||||
fsdp_mesh = self.device_mesh
|
||||
sharding_strategy = get_sharding_strategy(fsdp_mesh)
|
||||
|
||||
self.use_orig_params = fsdp_config.get("use_orig_params", False)
|
||||
if self.config.model.get("freeze_vision_tower", False):
|
||||
vision_tower = get_vl_model_vision_tower(critic_module)
|
||||
if vision_tower is not None:
|
||||
vision_tower.requires_grad_(False)
|
||||
self.use_orig_params = True
|
||||
if self.rank == 0:
|
||||
print("[critic model] Vision tower is set to not trainable.")
|
||||
else:
|
||||
if self.rank == 0:
|
||||
print("[critic model] No vision tower found.")
|
||||
|
||||
# Note: We force turn off CPUOffload for critic because it causes incorrect results when using grad accumulation
|
||||
if config.strategy == "fsdp":
|
||||
critic_module = FSDP(
|
||||
critic_module,
|
||||
param_init_fn=init_fn,
|
||||
use_orig_params=False,
|
||||
use_orig_params=self.use_orig_params,
|
||||
auto_wrap_policy=auto_wrap_policy,
|
||||
device_id=get_device_id(),
|
||||
sharding_strategy=sharding_strategy,
|
||||
|
||||
Reference in New Issue
Block a user