# VLM2Vec/src/trainer_add_CRD.py import collections import contextlib import functools import shutil import sys import time from datetime import timedelta from packaging import version from accelerate import skip_first_batches, DistributedType, InitProcessGroupKwargs from transformers import PretrainedConfig from transformers.trainer import Trainer, TRAINING_ARGS_NAME, TRAINER_STATE_NAME import torch.distributed as dist from typing import Optional import os import torch import math import torch.nn as nn from src.data.collator.train_collator import split_vlm_inputs, get_dense_rep, split_and_process_vlm_inputs from src.model.model_add_CRD import MMEBModel from src.loss_add_CRD import SimpleContrastiveLoss, DistributedContrastiveLoss, MultiLayerCRDLoss, DistributedMultiLayerCRDLoss from src.grad_cache.grad_cache import GradCache from torch.utils.data import DataLoader, Dataset, IterableDataset, RandomSampler, SequentialSampler from transformers.training_args import OptimizerNames, ParallelMode, TrainingArguments from transformers.trainer_callback import ( ExportableState, TrainerState, ) from transformers.trainer_utils import ( TrainOutput, has_length, speed_metrics, seed_worker, ) from transformers.trainer_pt_utils import ( get_model_param_count, ) from transformers.trainer import FSDP_MODEL_NAME from transformers.utils import ( XLA_FSDPV2_MIN_VERSION, is_accelerate_available, is_apex_available, is_torch_xla_available, logging, is_sagemaker_mp_enabled, CONFIG_NAME, WEIGHTS_NAME, SAFE_WEIGHTS_NAME, ADAPTER_WEIGHTS_NAME, ADAPTER_SAFE_WEIGHTS_NAME ) from src.utils import batch_to_device from src.utils import print_master, print_rank if is_apex_available(): from apex import amp if is_torch_xla_available(): import torch_xla.core.xla_model as xm from torch_xla import __version__ as XLA_VERSION IS_XLA_FSDPV2_POST_2_2 = version.parse(XLA_VERSION) >= version.parse(XLA_FSDPV2_MIN_VERSION) if IS_XLA_FSDPV2_POST_2_2: pass else: IS_XLA_FSDPV2_POST_2_2 = False logger = logging.get_logger(__name__) # =============== Helper: locate last LM block for grad diagnostics =============== def _locate_lm_layers_modulelist(encoder): """ Try common paths to get ModuleList of LM blocks in Qwen/LLaMA-like models. """ candidates = [ ("model", "language_model", "layers"), ("model", "model", "layers"), ("model", "layers"), ("language_model", "layers"), ("transformer", "layers"), ] for path in candidates: obj = encoder ok = True for p in path: if hasattr(obj, p): obj = getattr(obj, p) else: ok = False break if ok and isinstance(obj, torch.nn.ModuleList) and len(obj) > 0: return obj return None def _grad_norm(params): tot = 0.0 for p in params: if p.grad is not None: g = p.grad.detach().float() tot += (g * g).sum().item() return math.sqrt(tot) if tot > 0 else 0.0 # ================================================================================ class MMEBTrainer(Trainer): def __init__(self, *args, **kwargs): super(MMEBTrainer, self).__init__(*args, **kwargs) ws = dist.get_world_size() if dist.is_initialized() else 1 self.is_ddp = dist.is_initialized() and ws > 1 self._dist_loss_scale_factor = ws if self.is_ddp else 1 self.processor = self.processing_class def get_batch_samples(self, epoch_iterator, num_batches): batch_samples = [] num_items_in_batch = None for _ in range(num_batches): try: batch_samples += [next(epoch_iterator)] except StopIteration: break if len(batch_samples) > 0 and "labels" in batch_samples[0]: # For now we don't support object detection try: num_items_in_batch = sum([(batch["labels"].ne(-100)).sum() for batch in batch_samples]) except (TypeError, AttributeError): pass if self.args.average_tokens_across_devices and num_items_in_batch is not None: num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item() if torch.is_tensor(num_items_in_batch): num_items_in_batch = num_items_in_batch.item() return batch_samples, num_items_in_batch def compute_loss(self, model, inputs, *args, **kwargs): qry_inputs, tgt_inputs = inputs return model(qry=qry_inputs, tgt=tgt_inputs) def _save(self, output_dir: Optional[str] = None, state_dict=None): os.makedirs(output_dir, exist_ok=True) if state_dict is None: state_dict = self.model.state_dict() prefix = 'encoder.' assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys()) state_dict = {k[len(prefix):]: v for k, v in state_dict.items()} self.model.encoder.save_pretrained( output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors ) if self.tokenizer is not None: self.tokenizer.save_pretrained(output_dir) torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) def _get_train_sampler(self) -> Optional[torch.utils.data.Sampler]: # override original trainer's method if self.train_dataset is None or not has_length(self.train_dataset): return None return RandomSampler(self.train_dataset) def get_train_dataloader(self) -> DataLoader: """ override original trainer's method to disable self.accelerator.prepare since it will wrap DataLoaderDispatcher and lead to (1) `RuntimeError: You can't use batches of different size with `dispatch_batches=True` or when using an `IterableDataset`.` (2) all outputs of dataloader must be tensors """ if self.train_dataset is None: raise ValueError("Trainer: training requires a train_dataset.") train_dataset = self.train_dataset data_collator = self.data_collator train_dataset = self._remove_unused_columns(train_dataset, description="training") dataloader_params = { "batch_size": self._train_batch_size, "collate_fn": data_collator, "num_workers": self.args.dataloader_num_workers, "pin_memory": self.args.dataloader_pin_memory, "persistent_workers": self.args.dataloader_persistent_workers, } if not isinstance(train_dataset, torch.utils.data.IterableDataset): dataloader_params["sampler"] = self._get_train_sampler() dataloader_params["drop_last"] = self.args.dataloader_drop_last dataloader_params["worker_init_fn"] = seed_worker dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor else: dataloader_params["sampler"] = None dataloader_params["shuffle"] = False dataloader_params["drop_last"] = True dataloader_params["prefetch_factor"] = None # # tune on both prefetch_factor and persistent_workers will cause hang at epoch2 return DataLoader(train_dataset, **dataloader_params) def _load_from_checkpoint(self, resume_from_checkpoint, model=None): self.model_args.checkpoint_path = resume_from_checkpoint logger.info(f"Loading checkpoint from {resume_from_checkpoint}") self.model = MMEBModel.load(self.model_args) self.model_wrapped = self.model def _inner_training_loop( self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None ): self.accelerator.free_memory() self._train_batch_size = batch_size if self.args.auto_find_batch_size: if self.state.train_batch_size != self._train_batch_size: from accelerate.utils import release_memory (self.model_wrapped,) = release_memory(self.model_wrapped) self.model_wrapped = self.model # Check for DeepSpeed *after* the intial pass and modify the config if self.is_deepspeed_enabled: # Temporarily unset `self.args.train_batch_size` original_bs = self.args.per_device_train_batch_size self.args.per_device_train_batch_size = self._train_batch_size // max(1, self.args.n_gpu) self.propagate_args_to_deepspeed(True) self.args.per_device_train_batch_size = original_bs self.state.train_batch_size = self._train_batch_size logger.debug(f"Currently training with a batch size of: {self._train_batch_size}") # Data loader and number of training steps train_dataloader = self.get_train_dataloader() # Setting up training control variables: # number of training epochs: num_train_epochs # number of training steps per epoch: num_update_steps_per_epoch # total number of training steps to execute: max_steps total_train_batch_size = self._train_batch_size * args.gradient_accumulation_steps * args.world_size len_dataloader = None num_train_tokens = None if has_length(train_dataloader): len_dataloader = len(train_dataloader) num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1) num_examples = self.num_examples(train_dataloader) if args.max_steps > 0: max_steps = args.max_steps num_train_epochs = args.max_steps // num_update_steps_per_epoch + int( args.max_steps % num_update_steps_per_epoch > 0 ) # May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's # the best we can do. num_train_samples = args.max_steps * total_train_batch_size if args.include_tokens_per_second: num_train_tokens = ( self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps ) else: max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch) num_train_epochs = math.ceil(args.num_train_epochs) num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs if args.include_tokens_per_second: num_train_tokens = self.num_tokens(train_dataloader) * args.num_train_epochs elif args.max_steps > 0: # Rely on max_steps when dataloader does not have a working size max_steps = args.max_steps # Setting a very large number of epochs so we go as many times as necessary over the iterator. num_train_epochs = sys.maxsize num_update_steps_per_epoch = max_steps num_examples = total_train_batch_size * args.max_steps num_train_samples = args.max_steps * total_train_batch_size if args.include_tokens_per_second: num_train_tokens = self.num_tokens(train_dataloader, args.max_steps) * args.gradient_accumulation_steps else: raise ValueError( "args.max_steps must be set to a positive value if dataloader does not have a length, was" f" {args.max_steps}" ) delay_optimizer_creation = is_sagemaker_mp_enabled() or self.is_fsdp_xla_enabled or self.is_fsdp_enabled # We need to reset the scheduler, as its parameters may be different on subsequent calls if self._created_lr_scheduler: self.lr_scheduler = None self._created_lr_scheduler = False self.create_optimizer_and_scheduler(num_training_steps=max_steps) self.state = TrainerState( stateful_callbacks=[ cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState) ] ) self.state.is_hyper_param_search = trial is not None self.state.train_batch_size = self._train_batch_size # Compute absolute values for logging, eval, and save if given as ratio if args.logging_steps is not None: if args.logging_steps < 1: self.state.logging_steps = math.ceil(max_steps * args.logging_steps) else: self.state.logging_steps = args.logging_steps if args.eval_steps is not None: if args.eval_steps < 1: self.state.eval_steps = math.ceil(max_steps * args.eval_steps) else: self.state.eval_steps = args.eval_steps if args.save_steps is not None: if args.save_steps < 1: self.state.save_steps = math.ceil(max_steps * args.save_steps) else: self.state.save_steps = args.save_steps # Activate gradient checkpointing if needed if args.gradient_checkpointing: self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=args.gradient_checkpointing_kwargs) model = self._wrap_model(self.model_wrapped) # as the model is wrapped, don't use `accelerator.prepare` # this is for unhandled cases such as # FSDP-XLA, SageMaker MP/DP, DataParallel, IPEX use_accelerator_prepare = True if model is self.model else False if delay_optimizer_creation: if use_accelerator_prepare: self._fsdp_qlora_plugin_updates() self.model = self.accelerator.prepare(self.model) self.create_optimizer_and_scheduler(num_training_steps=max_steps) # prepare using `accelerator` prepare if use_accelerator_prepare: self.model.train() if hasattr(self.lr_scheduler, "step"): if self.use_apex: model = self.accelerator.prepare(self.model) else: model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer) else: # to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config. model, self.optimizer, self.lr_scheduler = self.accelerator.prepare( self.model, self.optimizer, self.lr_scheduler ) elif self.args.optim in [OptimizerNames.LOMO, OptimizerNames.ADALOMO]: # In this case we are in DDP + LOMO, which should be supported self.optimizer = self.accelerator.prepare(self.optimizer) if self.is_fsdp_enabled: self.model = self.model_wrapped = model # for the rest of this function `model` is the outside model, whether it was wrapped or not if model is not self.model: self.model_wrapped = model # backward compatibility if self.is_deepspeed_enabled: self.deepspeed = self.model_wrapped # Check if saved optimizer or scheduler states exist self._load_optimizer_and_scheduler(resume_from_checkpoint) # important: at this point: # self.model is the Transformers Model # self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), # FSDP(Transformers Model), Dynamo Optimized Module(Transformers Model) etc. # Train! logger.info("***** Running training *****") logger.info(f" Num examples = {num_examples:,}") logger.info(f" Num Epochs = {num_train_epochs:,}") logger.info(f" Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}") if self.args.per_device_train_batch_size != self._train_batch_size: logger.info(f" Training with DataParallel so batch size has been adjusted to: {self._train_batch_size:,}") logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size:,}") logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}") logger.info(f" Total optimization steps = {max_steps:,}") logger.info(f" Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}") self.state.epoch = 0 start_time = time.time() epochs_trained = 0 steps_trained_in_current_epoch = 0 steps_trained_progress_bar = None # @ruimeng use steps_trained_in_current_epoch to skip batches for finding buggy data # steps_trained_in_current_epoch = 42 # Check if continuing training from a checkpoint if resume_from_checkpoint is not None and os.path.isfile( os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME) ): self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)) self.compare_trainer_and_checkpoint_args(self.args, self.state) self._load_callback_state() epochs_trained = int(self.state.global_step // num_update_steps_per_epoch) if not args.ignore_data_skip: steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch) steps_trained_in_current_epoch *= args.gradient_accumulation_steps else: steps_trained_in_current_epoch = 0 logger.info(" Continuing training from checkpoint, will skip to saved global_step") logger.info(f" Continuing training from epoch {epochs_trained}") logger.info(f" Continuing training from global step {self.state.global_step}") if not args.ignore_data_skip: logger.info( f" Will skip the first {epochs_trained} epochs then the first" f" {steps_trained_in_current_epoch} batches in the first epoch." ) # Update the references self.callback_handler.model = self.model self.callback_handler.optimizer = self.optimizer self.callback_handler.lr_scheduler = self.lr_scheduler self.callback_handler.train_dataloader = train_dataloader # This should be the same if the state has been saved but in case the training arguments changed, it's safer # to set this after the load. self.state.max_steps = max_steps self.state.num_train_epochs = num_train_epochs self.state.is_local_process_zero = self.is_local_process_zero() self.state.is_world_process_zero = self.is_world_process_zero() # tr_loss is a tensor to avoid synchronization of TPUs through .item() tr_loss = torch.tensor(0.0).to(args.device) # _total_loss_scalar is updated everytime .item() has to be called on tr_loss and stores the sum of all losses self._total_loss_scalar = 0.0 self._globalstep_last_logged = self.state.global_step model.zero_grad() grad_norm: Optional[float] = None self.control = self.callback_handler.on_train_begin(args, self.state, self.control) if args.eval_on_start: self._evaluate(trial, ignore_keys_for_eval, skip_scheduler=True) total_batched_samples = 0 for epoch in range(epochs_trained, num_train_epochs): epoch_dataloader = train_dataloader if hasattr(epoch_dataloader.dataset, "set_epoch"): epoch_dataloader.dataset.set_epoch(epoch) if args.past_index >= 0: self._past = None steps_in_epoch = ( len(epoch_dataloader) if len_dataloader is not None else args.max_steps * args.gradient_accumulation_steps ) self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control) if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0: self._load_rng_state(resume_from_checkpoint) rng_to_sync = False steps_skipped = 0 if steps_trained_in_current_epoch > 0: epoch_dataloader = skip_first_batches(epoch_dataloader, steps_trained_in_current_epoch) steps_skipped = steps_trained_in_current_epoch steps_trained_in_current_epoch = 0 rng_to_sync = True step = -1 epoch_iterator = iter(epoch_dataloader) remainder = num_examples % args.gradient_accumulation_steps num_items_in_batch = None if remainder == 0: remainder = args.gradient_accumulation_steps update_step = -1 total_updates = steps_in_epoch // args.gradient_accumulation_steps + 1 for _ in range(total_updates): update_step += 1 num_batches = args.gradient_accumulation_steps if update_step != (total_updates - 1) else remainder batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches) for i, inputs in enumerate(batch_samples): step += 1 total_batched_samples += 1 dataset_stat = collections.Counter(inputs[0]['global_dataset_name']) is_last_step_and_steps_less_than_grad_acc = ( steps_in_epoch <= args.gradient_accumulation_steps and (step + 1) == steps_in_epoch ) do_sync_step = is_last_step_and_steps_less_than_grad_acc or ( total_batched_samples % args.gradient_accumulation_steps == 0 ) if not do_sync_step: self.accelerator.gradient_state._set_sync_gradients(False) else: self.accelerator.gradient_state._set_sync_gradients(True) if self.args.include_num_input_tokens_seen: main_input_name = getattr(self.model, "main_input_name", "input_ids") if main_input_name not in inputs: logger.warning("Tried to track the number of tokens seen, however the current model is not configured properly.") else: input_tokens = inputs[main_input_name].numel() input_tokens = torch.tensor(input_tokens, device=self.args.device, dtype=torch.int64) self.state.num_input_tokens_seen += self.accelerator.gather(input_tokens).cpu().item() if rng_to_sync: self._load_rng_state(resume_from_checkpoint) rng_to_sync = False if steps_trained_in_current_epoch > 0: steps_trained_in_current_epoch -= 1 if steps_trained_progress_bar is not None: steps_trained_progress_bar.update(1) if steps_trained_in_current_epoch == 0: self._load_rng_state(resume_from_checkpoint) continue elif steps_trained_progress_bar is not None: steps_trained_progress_bar.close() steps_trained_progress_bar = None if step % args.gradient_accumulation_steps == 0: self.control = self.callback_handler.on_step_begin(args, self.state, self.control) context = ( functools.partial(self.accelerator.no_sync, model=model) if i != len(batch_samples) - 1 else contextlib.nullcontext ) with context(): tr_loss_step = self.training_step(model, inputs, num_items_in_batch) if ( args.logging_nan_inf_filter and not is_torch_xla_available() and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step)) ): tr_loss = tr_loss + tr_loss / (1 + self.state.global_step - self._globalstep_last_logged) else: if tr_loss.device != tr_loss_step.device: raise ValueError(f"Calculated loss must be on the original device: {tr_loss.device} but device in use is {tr_loss_step.device}") tr_loss = tr_loss + tr_loss_step self.current_flos += float(self.floating_point_ops(inputs)) if do_sync_step: self.accelerator.gradient_state._set_sync_gradients(True) if args.max_grad_norm is not None and args.max_grad_norm > 0: if self.use_apex: _grad_norm = torch.nn.utils.clip_grad_norm_(amp.master_params(self.optimizer), args.max_grad_norm) else: _grad_norm = self.accelerator.clip_grad_norm_(model.parameters(), args.max_grad_norm) if (is_accelerate_available() and self.accelerator.distributed_type == DistributedType.DEEPSPEED): grad_norm = model.get_global_grad_norm() if hasattr(grad_norm, "item"): grad_norm = grad_norm.item() else: grad_norm = _grad_norm self.control = self.callback_handler.on_pre_optimizer_step(args, self.state, self.control) try: self._maybe_log_teacher_grad(model) except Exception as e: logger.warning(f"teacher grad log failed (ignored): {e}") self.optimizer.step() self.control = self.callback_handler.on_optimizer_step(args, self.state, self.control) optimizer_was_run = not self.accelerator.optimizer_step_was_skipped if optimizer_was_run: if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau): self.lr_scheduler.step() model.zero_grad() self.state.global_step += 1 self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch self.control = self.callback_handler.on_step_end(args, self.state, self.control) self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, time.time()) else: self.control = self.callback_handler.on_substep_end(args, self.state, self.control) if self.control.should_epoch_stop or self.control.should_training_stop: if is_torch_xla_available(): xm.mark_step() break if self.control.should_epoch_stop or self.control.should_training_stop: if is_torch_xla_available(): xm.mark_step() break if step < 0: logger.warning("There seems not to be a single sample in your epoch_iterator, stopping training.") self.control.should_training_stop = True self.control = self.callback_handler.on_epoch_end(args, self.state, self.control) self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval, time.time()) if self.control.should_training_stop: break if args.past_index and hasattr(self, "_past"): delattr(self, "_past") logger.info("\n\nTraining completed.\n\n") if args.load_best_model_at_end and self.state.best_model_checkpoint is not None: if is_torch_xla_available(): xm.rendezvous("load_best_model_at_end") elif args.parallel_mode == ParallelMode.DISTRIBUTED: dist.barrier() self._load_best_model() self._total_loss_scalar += tr_loss.item() effective_global_step = max(self.state.global_step, 0.001) train_loss = self._total_loss_scalar / effective_global_step metrics = speed_metrics( "train", start_time, num_samples=num_train_samples, num_steps=self.state.max_steps, num_tokens=num_train_tokens, ) self.store_flos() metrics["total_flos"] = self.state.total_flos metrics["train_loss"] = train_loss self.is_in_train = False self._memory_tracker.stop_and_update_metrics(metrics) self.log(metrics) run_dir = self._get_output_dir(trial) checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir) if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1: for checkpoint in checkpoints_sorted: if not os.path.samefile(checkpoint, self.state.best_model_checkpoint): logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit") shutil.rmtree(checkpoint, ignore_errors=True) self.control = self.callback_handler.on_train_end(args, self.state, self.control) self._finish_current_push() if self.neftune_noise_alpha is not None: self._deactivate_neftune(self.model) return TrainOutput(self.state.global_step, train_loss, metrics) class GradCacheLateProcessTrainer(MMEBTrainer): """ Adapted from gradcache repo. """ def __init__(self, *args, **kwargs): self.max_length = kwargs.get("max_length", 512) if "max_length" in kwargs: del kwargs["max_length"] self.model_args = kwargs.get("model_args", None) if "model_args" in kwargs: del kwargs["model_args"] super(GradCacheLateProcessTrainer, self).__init__(*args, **kwargs) ws = dist.get_world_size() if dist.is_initialized() else 1 self.is_ddp = dist.is_initialized() and ws > 1 self._dist_loss_scale_factor = ws if self.is_ddp else 1 loss_fn_cls = DistributedMultiLayerCRDLoss if self.is_ddp else MultiLayerCRDLoss crd_layers = getattr(self.args, "crd_layers", None) if isinstance(crd_layers, str) and len(crd_layers.strip()) > 0: crd_layers = [int(x.strip()) for x in crd_layers.split(",") if x.strip() != ""] else: crd_layers = None # ADDED: allow detach_teacher from args detach_teacher = getattr(self.args, "crd_detach_teacher", True) # 读取可选开关（也可以加到 TrainingArguments 后再改这里的读取） crd_side = os.getenv("CRD_SIDE", "both") # "both"|"qry"|"tgt" queue_size = int(os.getenv("CRD_QUEUE_SIZE", "0") or 0) self.loss_fn = loss_fn_cls( temperature=self.model.temperature, weights=getattr(self.model, "supervise_weights", None), crd_weight=getattr(self.args, "crd_weight", 0.2), crd_temperature=getattr(self.args, "crd_temperature", 0.07), crd_layers=crd_layers, detach_teacher=detach_teacher, crd_side=crd_side, queue_size=queue_size ) self.gc = GradCache( models=[self.model, self.model], chunk_sizes=[self.args.gc_q_chunk_size, self.args.gc_p_chunk_size], loss_fn=self.loss_fn, split_input_fn=split_and_process_vlm_inputs, get_rep_fn=get_dense_rep, # 返回 [B,K,D] fp16=self.args.fp16, scaler=self.scaler if self.args.fp16 else None ) # ADDED: cache for debug (last block params) self._last_block_params_cache = None def _infer_device(self, model): if hasattr(model, "device") and model.device is not None: return model.device try: return next(model.parameters()).device except StopIteration: return self.args.device def _batch_size(self, batch: dict) -> int: for k, v in batch.items(): if torch.is_tensor(v) and v.dim() > 0: return v.size(0) if isinstance(v, list): return len(v) raise ValueError("Cannot infer batch size from batch keys.") def _slice_batch(self, batch: dict, size: int, offset: int = 0): out = {} end = offset + size for k, v in batch.items(): if torch.is_tensor(v) and v.dim() > 0 and v.size(0) >= end: out[k] = v[offset:end] elif isinstance(v, list) and len(v) >= end: out[k] = v[offset:end] else: out[k] = v return out def _maybe_log_teacher_grad(self, model): """ 统计“最后一层 Block”的梯度范数，用于确认教师是否被更新。 可用环境变量 LOG_TEACHER_GRAD=0 关闭。 """ if os.getenv("LOG_TEACHER_GRAD", "1") not in ("1", "true", "True"): return try: params_last = self._get_last_block_params() tgn = _grad_norm(params_last) # 已转 float # 1) 控制台 print_master(f"[teacher_grad] step={self.state.global_step} norm={tgn:.6f}") # 2) HF metrics（会和 loss、grad_norm 一起进到 log_history/W&B） self.log({"teacher_grad_norm": tgn}) except Exception as e: logger.warning(f"teacher grad log failed: {e}") def _forward_reps_in_chunks(self, model, batch: dict, side: str, chunk_size: int, device): B = self._batch_size(batch) outs = [] use_bf16 = getattr(self.args, "bf16", False) dev_type = "cuda" if "cuda" in str(device) else "cpu" # 不切换 eval，不使用 no_grad；分块前向后立刻 detach for s in range(0, B, max(1, chunk_size or B)): bs = min(chunk_size or B, B - s) sub = self._slice_batch(batch, size=bs, offset=s) with torch.autocast(device_type=dev_type, dtype=torch.bfloat16, enabled=use_bf16): if side == "qry": o = model(qry=sub) reps = o["qry_reps"] # [bs,K,D] else: o = model(tgt=sub) reps = o["tgt_reps"] # [bs,K,D] outs.append(reps.detach()) # 关键：立即 detach del o, reps # 如显存非常紧张可偶尔清理缓存（频繁调用会慢） # if s % (4 * (chunk_size or B)) == 0: torch.cuda.empty_cache() return torch.cat(outs, dim=0) # [B,K,D] def _norm_weights(self, weights_list, K, device): if weights_list is None: return torch.ones(K, device=device) / K w = torch.tensor(list(weights_list), dtype=torch.float32, device=device) w = torch.clamp(w, min=0) s = float(w.sum()) return w / (s if s > 0 else 1.0) def _crd_indices(self, K: int): # 使用 args.crd_layers（相对 supervise_layers 的 0-based 索引）；None -> 所有非最后层 cl = getattr(self.args, "crd_layers", None) if cl is None or (isinstance(cl, str) and cl.strip() == ""): return list(range(0, max(0, K - 1))) if isinstance(cl, str): idxs = [int(x.strip()) for x in cl.split(",") if x.strip() != ""] else: idxs = list(cl) # 过滤掉最后一层 K-1 out = [] for i in idxs: if i < 0: i = K + i if 0 <= i < K - 1: out.append(i) return sorted(set(out)) def _single_gpu_chunked_step(self, model, queries: dict, targets: dict, device): """ 手写分块版：两阶段梯度 A) 对 query 分块，backward 到 q（y / teacher 常量） B) 对 target 分块，backward 到 p（x / teacher 常量） 返回一个不带梯度的标量 loss 供日志使用。 """ # 预先无梯度计算“常量库”：全 batch 的多层表示（占用小） q_all = self._forward_reps_in_chunks(model, queries, side="qry", chunk_size=self.args.gc_q_chunk_size, device=device) # [B,K,D] p_all = self._forward_reps_in_chunks(model, targets, side="tgt", chunk_size=self.args.gc_p_chunk_size, device=device) # [B,K,D] B, K, D = q_all.shape w_ret = self._norm_weights(getattr(self.model, "supervise_weights", None), K, device) crd_idxs = self._crd_indices(K) w_crd = self._norm_weights([w_ret[i].item() for i in crd_idxs] if len(crd_idxs) > 0 else [1.0], max(1, len(crd_idxs)), device) temp = float(getattr(self.model, "temperature", 0.02)) beta = float(getattr(self.loss_fn, "runtime_beta", getattr(self.loss_fn, "crd_weight", getattr(self.args, "crd_weight", 0.2)))) crd_temp = float(getattr(self.loss_fn, "crd_temperature", getattr(self.args, "crd_temperature", 0.07))) detach_teacher = bool(getattr(self.args, "crd_detach_teacher", True)) use_bf16 = getattr(self.args, "bf16", False) dev_type = "cuda" if "cuda" in str(device) else "cpu" # 教师（最后一层）常量 tq_all = q_all[:, K - 1, :].detach() if detach_teacher else q_all[:, K - 1, :] tp_all = p_all[:, K - 1, :].detach() if detach_teacher else p_all[:, K - 1, :] total_loss_scalar = 0.0 # Phase-A: 查询侧分块，更新 q q_chunk = max(1, self.args.gc_q_chunk_size or B) for s in range(0, B, q_chunk): bs = min(q_chunk, B - s) sub_q = self._slice_batch(queries, size=bs, offset=s) labels = torch.arange(s, s + bs, device=device, dtype=torch.long) with torch.autocast(device_type=dev_type, dtype=torch.bfloat16, enabled=use_bf16): o_q = model(qry=sub_q) # [bs,K,D]，带梯度 qk = o_q["qry_reps"] # L_ret (q_k vs p_all_k) L_ret_q = 0.0 for k_idx in range(K): logits = torch.matmul(qk[:, k_idx, :], p_all[:, k_idx, :].transpose(0, 1)) / temp # [bs,B] Lk = torch.nn.functional.cross_entropy(logits, labels, reduction="mean") L_ret_q = L_ret_q + w_ret[k_idx] * Lk # L_crd_q (q_k vs tq_all) L_crd_q = 0.0 for j, k_idx in enumerate(crd_idxs): logits = torch.matmul(qk[:, k_idx, :], tq_all.transpose(0, 1)) / crd_temp Lk = torch.nn.functional.cross_entropy(logits, labels, reduction="mean") L_crd_q = L_crd_q + w_crd[j] * Lk L_q = L_ret_q + beta * L_crd_q # 反传仅更新 q 分支 L_q.backward() total_loss_scalar += float(L_q.detach()) * (bs / B) del qk, o_q, L_q, L_ret_q, L_crd_q # Phase-B: 候选侧分块，更新 p p_chunk = max(1, self.args.gc_p_chunk_size or B) for s in range(0, B, p_chunk): bs = min(p_chunk, B - s) sub_p = self._slice_batch(targets, size=bs, offset=s) labels = torch.arange(s, s + bs, device=device, dtype=torch.long) with torch.autocast(device_type=dev_type, dtype=torch.bfloat16, enabled=use_bf16): o_p = model(tgt=sub_p) # [bs,K,D]，带梯度 pk = o_p["tgt_reps"] # L_ret (p_k vs q_all_k) L_ret_p = 0.0 for k_idx in range(K): logits = torch.matmul(pk[:, k_idx, :], q_all[:, k_idx, :].transpose(0, 1)) / temp Lk = torch.nn.functional.cross_entropy(logits, labels, reduction="mean") L_ret_p = L_ret_p + w_ret[k_idx] * Lk # L_crd_p (p_k vs tp_all) L_crd_p = 0.0 for j, k_idx in enumerate(crd_idxs): logits = torch.matmul(pk[:, k_idx, :], tp_all.transpose(0, 1)) / crd_temp Lk = torch.nn.functional.cross_entropy(logits, labels, reduction="mean") L_crd_p = L_crd_p + w_crd[j] * Lk L_p = L_ret_p + beta * L_crd_p L_p.backward() total_loss_scalar += float(L_p.detach()) * (bs / B) del pk, o_p, L_p, L_ret_p, L_crd_p # 返回常数张量用于日志（避免 HF 再做 backward） return torch.tensor(total_loss_scalar, device=device, dtype=torch.float32, requires_grad=False) # ADDED: dynamic beta warmup and grad diagnostics def _apply_crd_warmup(self): target_beta = getattr(self.args, "crd_weight", 0.2) warm_steps = getattr(self.args, "crd_warmup_steps", 0) # self.state.global_step 在 accumulation 前是上一次的值，用它也足够 step = max(0, getattr(self.state, "global_step", 0)) if warm_steps and step < warm_steps: beta = target_beta * float(step + 1) / float(warm_steps) else: beta = target_beta # 同时写入 runtime_beta 与 crd_weight，保证兼容不同实现 setattr(self.loss_fn, "runtime_beta", beta) if hasattr(self.loss_fn, "crd_weight"): self.loss_fn.crd_weight = beta def _get_last_block_params(self): if self._last_block_params_cache is not None: return self._last_block_params_cache layers = _locate_lm_layers_modulelist(self.model.encoder) if layers is None or len(layers) == 0: logger.warning("Could not locate LM layers for grad debug; will use encoder parameters as fallback.") params = list(self.model.encoder.parameters()) else: params = list(layers[-1].parameters()) self._last_block_params_cache = params return params def _debug_teacher_grad(self, queries, targets, model): # 单卡禁用，避免两次整批前向导致 OOM if not self.is_ddp: return dbg_every = int(getattr(self.args, "crd_debug_every", 0) or 0) if dbg_every <= 0 or (self.state.global_step % dbg_every) != 0: return last_params = self._get_last_block_params() beta_saved = getattr(self.loss_fn, "crd_weight", 0.0) if self.is_ddp: # DDP: 两次经过 GradCache self.model.zero_grad(set_to_none=True) self.loss_fn.crd_weight = 0.0 _ = self.gc(queries, targets, no_sync_except_last=True) gn_ret = _grad_norm(last_params) self.model.zero_grad(set_to_none=True) self.loss_fn.crd_weight = beta_saved _ = self.gc(queries, targets, no_sync_except_last=True) gn_all = _grad_norm(last_params) else: # 单卡: 直接模型前向 + loss_fn self.model.zero_grad(set_to_none=True) # 取多层表征 out_q = model(qry=queries["qry"]) out_p = model(tgt=targets["tgt"]) x, y = out_q["qry_reps"], out_p["tgt_reps"] self.loss_fn.crd_weight = 0.0 loss = self.loss_fn(x, y) loss.backward() gn_ret = _grad_norm(last_params) self.model.zero_grad(set_to_none=True) self.loss_fn.crd_weight = beta_saved loss = self.loss_fn(x, y) loss.backward() gn_all = _grad_norm(last_params) print_master(f"[CRD-Debug] step={self.state.global_step} grad_norm(last-block): RET={gn_ret:.6f}, RET+CRD={gn_all:.6f}, delta={max(0.0, gn_all-gn_ret):.6f}") self.model.zero_grad(set_to_none=True) def training_step(self, model, inputs, *args, **kwargs) -> torch.Tensor: model.train() queries, targets = inputs device = self._infer_device(model) queries = batch_to_device(queries, device) targets = batch_to_device(targets, device) queries, targets = {'qry': queries}, {'tgt': targets} # 动态 CRD warmup self._apply_crd_warmup() # 可选梯度诊断 try: self._debug_teacher_grad(queries, targets, model) except Exception as e: logger.warning(f"CRD grad debug failed (ignored): {e}") if self.is_ddp: # 多卡：使用 GradCache（要求模型已被 DDP 包裹，HF 会在 _wrap_model + accelerator.prepare 后处理） self.gc.models = [model, model] loss = self.gc(queries, targets, no_sync_except_last=True) else: # 单卡：手写分块两阶段，避免整批前向 OOM loss = self._single_gpu_chunked_step(model, queries["qry"], targets["tgt"], device) return loss / self._dist_loss_scale_factor def _save(self, output_dir: Optional[str] = None, state_dict=None): print_master(f"Saving model to {output_dir}") os.makedirs(output_dir, exist_ok=True) if state_dict is None: state_dict = self.model.state_dict() prefix = 'encoder.' assert all(k.startswith(prefix) for k in state_dict.keys()), list(state_dict.keys()) state_dict = {k[len(prefix):]: v for k, v in state_dict.items()} self.model.encoder.save_pretrained( output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors ) if self.tokenizer is not None: self.tokenizer.save_pretrained(output_dir) torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) self.model.encoder.config.to_json_file(os.path.join(output_dir, 'config.json'))