wehub-resource-sync
193 lines
7.6 KiB
Python
193 lines
7.6 KiB
Python
# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
|
|
from contextlib import nullcontext
|
|
|
|
import paddle
|
|
from paddle.base import core
|
|
from paddle.distributed import fleet
|
|
|
|
from paddlenlp.trainer import Trainer
|
|
from paddlenlp.transformers.contrastive_loss import (
|
|
MatryoshkaContrastiveLoss,
|
|
MatryoshkaInfclLoss,
|
|
SimpleContrastiveLoss,
|
|
SimpleInfclLoss,
|
|
)
|
|
from paddlenlp.transformers.embedding_utils import dist_gather_tensor_with_gradient
|
|
from paddlenlp.utils import empty_device_cache
|
|
|
|
__all__ = ["EmbeddingTrainer"]
|
|
|
|
|
|
class EmbeddingTrainer(Trainer):
|
|
def __init__(self, model_args, **kwargs):
|
|
super().__init__(**kwargs)
|
|
|
|
self.model_args = model_args
|
|
self.embedding_negatives_cross_device = model_args.embedding_negatives_cross_device
|
|
self.accum_data = []
|
|
self.accum_freq = 0
|
|
self.accum_q_features = []
|
|
self.accum_p_features = []
|
|
self.accum_rng_states = {}
|
|
self.accum_rng_states["cpu"] = []
|
|
self.accum_rng_states["cuda"] = []
|
|
self.accum_rng_states["hybrid"] = []
|
|
|
|
if model_args.embedding_matryoshka_dims is not None and len(model_args.embedding_matryoshka_dims) > 0:
|
|
if model_args.loss_type == "inf_cl":
|
|
self.embedding_negatives_cross_device = False
|
|
self.loss_fn = MatryoshkaInfclLoss(model_args.embedding_matryoshka_dims, model_args.inf_cl_head_dim)
|
|
elif model_args.loss_type == "contrastive":
|
|
self.loss_fn = MatryoshkaContrastiveLoss(
|
|
model_args.embedding_temperature, model_args.embedding_matryoshka_dims
|
|
)
|
|
else:
|
|
if model_args.loss_type == "inf_cl":
|
|
self.embedding_negatives_cross_device = False
|
|
self.loss_fn = SimpleInfclLoss(model_args.inf_cl_head_dim)
|
|
elif model_args.loss_type == "contrastive":
|
|
self.loss_fn = SimpleContrastiveLoss(model_args.embedding_temperature)
|
|
|
|
def clear_memory(self):
|
|
self.accum_q_features.clear()
|
|
self.accum_p_features.clear()
|
|
empty_device_cache()
|
|
|
|
def clear_state(self):
|
|
self.accum_data.clear()
|
|
self.accum_rng_states["cpu"].clear()
|
|
self.accum_rng_states["cuda"].clear()
|
|
self.accum_rng_states["hybrid"].clear()
|
|
self.accum_freq = 0
|
|
|
|
@paddle.no_grad()
|
|
def forward_no_grad(self, model, inputs):
|
|
# Step1: graph-less forward
|
|
self.accum_data.append(inputs)
|
|
inputs = self._prepare_inputs(inputs)
|
|
with self.autocast_smart_context_manager():
|
|
# collect rand states
|
|
self.accum_rng_states["cpu"].append(paddle.framework.core.default_cpu_generator().get_state())
|
|
self.accum_rng_states["cuda"].append(paddle.get_rng_state())
|
|
if self.args.use_hybrid_parallel:
|
|
self.accum_rng_states["hybrid"].append(
|
|
fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()
|
|
)
|
|
|
|
query_reps, passage_reps = model(**inputs, return_encode=True)
|
|
|
|
if self.embedding_negatives_cross_device:
|
|
query_reps = dist_gather_tensor_with_gradient(query_reps)
|
|
passage_reps = dist_gather_tensor_with_gradient(passage_reps)
|
|
|
|
self.accum_q_features.append(query_reps)
|
|
self.accum_p_features.append(passage_reps)
|
|
|
|
self.accum_freq += 1
|
|
|
|
def get_current_rng_state(self):
|
|
return {
|
|
"cpu": [paddle.framework.core.default_cpu_generator().get_state()],
|
|
"cuda": [paddle.get_rng_state()],
|
|
"hybrid": [fleet.meta_parallel.get_rng_state_tracker().get_states_tracker()]
|
|
if self.args.use_hybrid_parallel
|
|
else [],
|
|
}
|
|
|
|
def reset_rng_state(self, states, index=0):
|
|
# set random states
|
|
if len(states) != 3:
|
|
raise ValueError("The length of state should be 3")
|
|
cpu_state = states["cpu"][index]
|
|
cuda_state = states["cuda"][index]
|
|
paddle.framework.core.default_cpu_generator().set_state(cpu_state)
|
|
# TODO(daisiming): support xpu and other custom devices.
|
|
if core.is_compiled_with_cuda():
|
|
for j in range(core.get_cuda_device_count()):
|
|
core.default_cuda_generator(j).set_state(cuda_state[j])
|
|
if self.args.use_hybrid_parallel:
|
|
hybrid_state = states["hybrid"][index]
|
|
fleet.meta_parallel.get_rng_state_tracker().set_states_tracker(hybrid_state)
|
|
|
|
def accum_forward_backward(self, model):
|
|
# Step2: representation gradient computation and caching
|
|
for i in range(len(self.accum_q_features)):
|
|
self.accum_q_features[i].stop_gradient = False
|
|
q_reps = paddle.concat(self.accum_q_features, axis=0)
|
|
for i in range(len(self.accum_p_features)):
|
|
self.accum_p_features[i].stop_gradient = False
|
|
p_reps = paddle.concat(self.accum_p_features, axis=0)
|
|
|
|
loss = self.loss_fn(q_reps, p_reps)
|
|
if self.do_grad_scaling:
|
|
self.scaler.scale(loss).backward()
|
|
else:
|
|
loss.backward()
|
|
# get representation gradient cache
|
|
accum_q_grads = [q.grad for q in self.accum_q_features]
|
|
accum_p_grads = [p.grad for p in self.accum_p_features]
|
|
del q_reps, p_reps
|
|
|
|
# clear trash memory
|
|
self.clear_memory()
|
|
|
|
current_rng_state = self.get_current_rng_state()
|
|
# Step3: sub-batch gradient accumulation
|
|
for i in range(self.accum_freq):
|
|
inputs = self.accum_data[i]
|
|
inputs = self._prepare_inputs(inputs)
|
|
|
|
sync_context = model.no_sync() if i != self.accum_freq - 1 and hasattr(model, "no_sync") else nullcontext()
|
|
with sync_context:
|
|
self.reset_rng_state(self.accum_rng_states, index=i)
|
|
|
|
with self.autocast_smart_context_manager():
|
|
query_reps, passage_reps = model(**inputs, return_encode=True)
|
|
|
|
if self.embedding_negatives_cross_device:
|
|
query_reps = dist_gather_tensor_with_gradient(query_reps)
|
|
passage_reps = dist_gather_tensor_with_gradient(passage_reps)
|
|
|
|
_loss = paddle.dot(query_reps.flatten(), accum_q_grads[i].flatten()) + paddle.dot(
|
|
passage_reps.flatten(), accum_p_grads[i].flatten()
|
|
)
|
|
_loss.backward()
|
|
|
|
self.reset_rng_state(current_rng_state)
|
|
self.clear_state()
|
|
return loss.detach()
|
|
|
|
def training_step(
|
|
self,
|
|
model,
|
|
inputs,
|
|
step_control=0,
|
|
):
|
|
if self.args.pipeline_parallel_degree > 1:
|
|
raise NotImplementedError("Cannot support pipeline parallel for Embedding training now.")
|
|
|
|
if self.args.gradient_accumulation_steps == 1:
|
|
return super().training_step(model, inputs)
|
|
else:
|
|
self.forward_no_grad(model, inputs)
|
|
|
|
# if (step_control + 1) % self.args.gradient_accumulation_steps is not zero, move on to next batch.
|
|
if (step_control + 1) % self.args.gradient_accumulation_steps != 0:
|
|
return 0.0
|
|
|
|
loss = self.accum_forward_backward(model)
|
|
return loss
|