diff --git a/README.md b/README.md index 036ee648..77d9c709 100644 --- a/README.md +++ b/README.md @@ -298,7 +298,7 @@ huggingface-cli login | datasets | 2.16.0 | 2.19.2 | | accelerate | 0.30.1 | 0.30.1 | | peft | 0.11.1 | 0.11.1 | -| trl | 0.8.6 | 0.9.3 | +| trl | 0.8.6 | 0.9.4 | | Optional | Minimum | Recommend | | ------------ | ------- | --------- | diff --git a/README_zh.md b/README_zh.md index fb0d790e..9a52a963 100644 --- a/README_zh.md +++ b/README_zh.md @@ -298,7 +298,7 @@ huggingface-cli login | datasets | 2.16.0 | 2.19.2 | | accelerate | 0.30.1 | 0.30.1 | | peft | 0.11.1 | 0.11.1 | -| trl | 0.8.6 | 0.9.3 | +| trl | 0.8.6 | 0.9.4 | | 可选项 | 至少 | 推荐 | | ------------ | ------- | --------- | diff --git a/src/llamafactory/train/ppo/ppo_utils.py b/src/llamafactory/train/ppo/ppo_utils.py index 570409f2..cfda0b3d 100644 --- a/src/llamafactory/train/ppo/ppo_utils.py +++ b/src/llamafactory/train/ppo/ppo_utils.py @@ -1,7 +1,9 @@ import json +from contextlib import nullcontext from typing import TYPE_CHECKING, Dict, List, Literal, Optional import torch +from transformers.integrations import is_deepspeed_zero3_enabled from ...extras.packages import is_requests_available @@ -28,16 +30,27 @@ def get_rewards_from_server(server_url: str, messages: List[str]) -> List[torch. def replace_model(model: "AutoModelForCausalLMWithValueHead", target: Literal["default", "reward"]) -> None: r""" - Replaces the default/reward modules in the model. The model is already unwrapped (and gathered). + Replaces the default/reward modules in the model. The model is already unwrapped. """ - if target == "reward": # save default head temporarily - setattr(model, "default_head_weight", model.v_head.summary.weight.data.detach().clone()) - setattr(model, "default_head_bias", model.v_head.summary.bias.data.detach().clone()) + if is_deepspeed_zero3_enabled(): + import deepspeed # type: ignore + + params = [model.v_head.summary.weight, model.v_head.summary.bias] + context_maybe_zero3 = deepspeed.zero.GatheredParameters(params, modifier_rank=0) + else: + context_maybe_zero3 = nullcontext() model.pretrained_model.set_adapter(target) # set the LoRA adapter to be active - device = model.v_head.summary.weight.device - model.v_head.summary.weight.data = model.get_buffer("{}_head_weight".format(target)).detach().clone().to(device) - model.v_head.summary.bias.data = model.get_buffer("{}_head_bias".format(target)).detach().clone().to(device) + with context_maybe_zero3: + if target == "reward": # save default head temporarily + setattr(model, "default_head_weight", model.v_head.summary.weight.data.detach().clone()) + setattr(model, "default_head_bias", model.v_head.summary.bias.data.detach().clone()) + + device = model.v_head.summary.weight.device + model.v_head.summary.weight.data = ( + model.get_buffer("{}_head_weight".format(target)).detach().clone().to(device) + ) + model.v_head.summary.bias.data = model.get_buffer("{}_head_bias".format(target)).detach().clone().to(device) def dump_layernorm(model: "PreTrainedModel") -> Dict[str, torch.Tensor]: diff --git a/src/llamafactory/train/ppo/trainer.py b/src/llamafactory/train/ppo/trainer.py index 7addfc3c..0ad1b8e8 100644 --- a/src/llamafactory/train/ppo/trainer.py +++ b/src/llamafactory/train/ppo/trainer.py @@ -309,12 +309,6 @@ class CustomPPOTrainer(PPOTrainer, Trainer): ) return lr_scheduler - def _save(self, output_dir: Optional[str] = None, state_dict: Optional[Dict[str, "torch.Tensor"]] = None) -> None: - super()._save(output_dir, state_dict) - if self.processor is not None: - output_dir = output_dir if output_dir is not None else self.args.output_dir - getattr(self.processor, "image_processor").save_pretrained(output_dir) - @torch.no_grad() def get_inputs(self, batch: Dict[str, "torch.Tensor"]) -> Tuple[List["torch.Tensor"], List["torch.Tensor"]]: r""" @@ -326,6 +320,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer): batch[k] = v[:, start_index:] with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model: + unwrapped_model = self.accelerator.unwrap_model(self.model) # issue in trl v0.8.6 if self.model_args.upcast_layernorm: layernorm_params = dump_layernorm(unwrapped_model) @@ -369,19 +364,19 @@ class CustomPPOTrainer(PPOTrainer, Trainer): return get_rewards_from_server(self.reward_model, messages) batch = self.prepare_model_inputs(queries, responses) + unwrapped_model: "AutoModelForCausalLMWithValueHead" = self.accelerator.unwrap_model(self.model) - with unwrap_model_for_generation(self.model, self.accelerator) as unwrapped_model: - if self.finetuning_args.reward_model_type == "lora": - replace_model(unwrapped_model, target="reward") - reward_model = self.model - else: - reward_model = self.reward_model + if self.finetuning_args.reward_model_type == "lora": + replace_model(unwrapped_model, target="reward") + reward_model = self.model + else: + reward_model = self.reward_model - with self.amp_context: # support bf16 - _, _, values = reward_model(**batch, output_hidden_states=True, return_dict=True, use_cache=False) + with unwrap_model_for_generation(reward_model, self.accelerator), self.amp_context: # support bf16 + _, _, values = reward_model(**batch, output_hidden_states=True, return_dict=True, use_cache=False) - if self.finetuning_args.reward_model_type == "lora": - replace_model(unwrapped_model, target="default") + if self.finetuning_args.reward_model_type == "lora": + replace_model(unwrapped_model, target="default") if self.is_chatglm_model: # assume same architecture values = torch.transpose(values, 0, 1) @@ -482,3 +477,7 @@ class CustomPPOTrainer(PPOTrainer, Trainer): self._save(output_dir, state_dict={}) remove_dummy_checkpoint(True, output_dir, [WEIGHTS_NAME, SAFE_WEIGHTS_NAME]) self.model.save_checkpoint(output_dir) + + if self.processor is not None: + output_dir = output_dir if output_dir is not None else self.args.output_dir + getattr(self.processor, "image_processor").save_pretrained(output_dir)