HuggingFace Trainer 源码解析：权重保存与模型检查点机制

2. control.should_log 与 is_torch_tpu_available

self.control.should_log 来源 self.control，is_torch_tpu_available 是判断 torch_xla 安装情况。若安装则使用 TPU 环境条件。

if self.control.should_log:
    if is_torch_tpu_available():
        xm.mark_step()

3. 评估（self.control.should_evaluate）

需要评估方法，得到评估值 metrics。

metrics = None
if self.control.should_evaluate:
    if isinstance(self.eval_dataset, dict):
        metrics = {}
        for eval_dataset_name, eval_dataset in self.eval_dataset.items():
            dataset_metrics = self.evaluate(
                eval_dataset=eval_dataset,
                ignore_keys=ignore_keys_for_eval,
                metric_key_prefix=f"eval_{eval_dataset_name}",
            )
            metrics.update(dataset_metrics)
    else:
        metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
    self._report_to_hp_search(trial, self.state.global_step, metrics)
    # Run delayed LR scheduler now that metrics are populated
    if isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
        metric_to_check = self.args.metric_for_best_model
        if not metric_to_check.startswith("eval_"):
            metric_to_check = f"eval_{metric_to_check}"
        self.lr_scheduler.step(metrics[metric_to_check])

4. 保存权重

self.control.should_save 由 self.control 给出，与权重保存相关的是 self._save_checkpoint 函数调用。

if self.control.should_save:
    self._save_checkpoint(model, trial, metrics=metrics)
    self.control = self.callback_handler.on_save(self.args, self.state, self.control)

5. self.callback_handler 类

self.callback_handler 负责调用回调列表。其大致逻辑如下：

class CallbackHandler(TrainerCallback):
    """Internal class that just calls the list of callbacks in order."""
    def __init__(self, callbacks, model, tokenizer, optimizer, lr_scheduler):
        self.callbacks = []
        for cb in callbacks:
            self.add_callback(cb)
        self.model = model
        self.tokenizer = tokenizer
        self.optimizer = optimizer
        self.lr_scheduler = lr_scheduler
        self.train_dataloader = None
        self.eval_dataloader = None
        # ... initialization logic ...

    def call_event(self, event, args, state, control, **kwargs):
        for callback in self.callbacks:
            result = getattr(callback, event)(
                args, state, control, model=self.model, tokenizer=self.tokenizer,
                optimizer=self.optimizer, lr_scheduler=self.lr_scheduler,
                train_dataloader=self.train_dataloader, eval_dataloader=self.eval_dataloader,
                **kwargs,
            )
            if result is not None:
                control = result
        return control

三、self._save_checkpoint 源码解读

1. 完整的源码

def _save_checkpoint(self, model, trial, metrics=None):
    # In all cases, including ddp/dp/deepspeed, self.model is always a reference to the model we want to save except FullyShardedDDP.
    assert unwrap_model(model) is self.model, "internal model should be a reference to self.model"
    
    # Save model checkpoint
    checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
    if self.hp_search_backend is None and trial is None:
        self.store_flos()
    run_dir = self._get_output_dir(trial=trial)
    output_dir = os.path.join(run_dir, checkpoint_folder)
    self.save_model(output_dir, _internal_call=True)
    
    if self.is_deepspeed_enabled:
        # under zero3 model file itself doesn't get saved since it's bogus!
        self.model_wrapped.save_checkpoint(output_dir)
    
    # Save optimizer and scheduler
    if self.sharded_ddp == ShardedDDPOption.SIMPLE:
        self.optimizer.consolidate_state_dict()
    
    if self.fsdp or self.is_fsdp_enabled:
        if self.is_fsdp_enabled:
            save_fsdp_optimizer(
                self.accelerator.state.fsdp_plugin, self.accelerator, self.optimizer, self.model, output_dir
            )
        else:
            full_osd = self.model.__class__.full_optim_state_dict(self.model, self.optimizer)
    
    if is_torch_tpu_available():
        xm.rendezvous("saving_optimizer_states")
        xm.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
        with warnings.catch_warnings(record=True) as caught_warnings:
            xm.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
            reissue_pt_warnings(caught_warnings)
    elif is_sagemaker_mp_enabled():
        opt_state_dict = self.optimizer.local_state_dict(gather_if_shard=False)
        smp.barrier()
        if smp.rdp_rank() == 0 or smp.state.cfg.shard_optimizer_state:
            smp.save(opt_state_dict, os.path.join(output_dir, OPTIMIZER_NAME), partial=True, v3=smp.state.cfg.shard_optimizer_state)
        if self.args.should_save:
            with warnings.catch_warnings(record=True) as caught_warnings:
                torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
                reissue_pt_warnings(caught_warnings)
            if self.do_grad_scaling:
                torch.save(self.scaler.state_dict(), os.path.join(output_dir, SCALER_NAME))
    elif self.args.should_save and not self.is_deepspeed_enabled:
        if self.fsdp and not self.is_fsdp_enabled:
            torch.save(full_osd, os.path.join(output_dir, OPTIMIZER_NAME))
        else:
            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
        with warnings.catch_warnings(record=True) as caught_warnings:
            torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, SCHEDULER_NAME))
            reissue_pt_warnings(caught_warnings)
        if self.do_grad_scaling:
            torch.save(self.scaler.state_dict(), os.path.join(output_dir, SCALER_NAME))
    
    # Determine the new best metric / best model checkpoint
    if metrics is not None and self.args.metric_for_best_model is not None:
        metric_to_check = self.args.metric_for_best_model
        if not metric_to_check.startswith("eval_"):
            metric_to_check = f"eval_{metric_to_check}"
        metric_value = metrics[metric_to_check]
        operator = np.greater if self.args.greater_is_better else np.less
        if (self.state.best_metric is None or self.state.best_model_checkpoint is None or operator(metric_value, self.state.best_metric)):
            self.state.best_metric = metric_value
            self.state.best_model_checkpoint = output_dir
    
    # Save the Trainer state
    if self.args.should_save:
        self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))
    
    # Save RNG state in non-distributed training
    rng_states = {
        "python": random.getstate(),
        "numpy": np.random.get_state(),
        "cpu": torch.random.get_rng_state(),
    }
    if torch.cuda.is_available():
        if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
            rng_states["cuda"] = torch.cuda.random.get_rng_state_all()
        else:
            rng_states["cuda"] = torch.cuda.random.get_rng_state()
    if is_torch_tpu_available():
        rng_states["xla"] = xm.get_rng_state()
    
    os.makedirs(output_dir, exist_ok=True)
    if self.args.world_size <= 1:
        torch.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
    else:
        torch.save(rng_states, os.path.join(output_dir, f"rng_state_{self.args.process_index}.pth"))
    
    if self.args.push_to_hub:
        self._push_from_checkpoint(output_dir)
    
    # Maybe delete some older checkpoints.
    if self.args.should_save:
        self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)

2. 获得保存路径与模型参数计算

output_dir 为最终保存路径（如：./out_dirs/checkpoint-1），self.store_flos() 为模型参数计算。

checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
if self.hp_search_backend is None and trial is None:
    self.store_flos()
run_dir = self._get_output_dir(trial=trial)
output_dir = os.path.join(run_dir, checkpoint_folder)

self.store_flos() 参数计算

def store_flos(self):
    # Storing the number of floating-point operations that went into the model
    if self.args.parallel_mode == ParallelMode.DISTRIBUTED:
        self.state.total_flos += (
            distributed_broadcast_scalars([self.current_flos], device=self.args.device).sum().item()
        )
        self.current_flos = 0
    else:
        self.state.total_flos += self.current_flos
        self.current_flos = 0

3. self.save_model(output_dir, _internal_call=True) 源码解读

该方法根据不同的条件和设置，选择性地保存模型的权重参数。涵盖了对不同分布式训练策略（如 DDP、DeepSpeed 等）的处理逻辑。

self.save_model(output_dir, _internal_call=True)

4. self.sharded_ddp == ShardedDDPOption.SIMPLE 条件保存 optimizer and scheduler 状态

如果使用的分布式训练策略是 ShardedDDPOption.SIMPLE，则调用了 optimizer 对象的 consolidate_state_dict() 方法，将优化器的状态字典整合成一个单独的状态字典。

if self.sharded_ddp == ShardedDDPOption.SIMPLE:
    self.optimizer.consolidate_state_dict()

5. self.fsdp or self.is_fsdp_enabled 条件保存优化器状态

在使用 FullyShardedDDP（FSDP）分布式训练策略时，根据 FSDP 的状态和设置来选择合适的方式保存优化器状态。

if self.fsdp or self.is_fsdp_enabled:
    if self.is_fsdp_enabled:
        save_fsdp_optimizer(
            self.accelerator.state.fsdp_plugin, self.accelerator, self.optimizer, self.model, output_dir
        )
    else:
        full_osd = self.model.__class__.full_optim_state_dict(self.model, self.optimizer)

6. 不同条件保存 optimizer、scheduler and scaler 等状态

根据训练环境和设置，选择性地保存优化器和学习率调度器的状态。包括 Torch TPU、SageMaker Model Parallelism 以及标准 PyTorch 环境下的保存逻辑。

if is_torch_tpu_available():
    xm.rendezvous("saving_optimizer_states")
    xm.save(self.optimizer.state_dict(), os.path.join(output_dir, OPTIMIZER_NAME))
    # ... save scheduler ...
elif is_sagemaker_mp_enabled():
    # ... save optimizer ...
    # ... save scheduler ...
elif self.args.should_save and not self.is_deepspeed_enabled:
    # ... save optimizer ...
    # ... save scheduler ...

其中文件名为：OPTIMIZER_NAME='optimizer.pt', SCHEDULER_NAME='scheduler.pt', SCALER_NAME='scaler.pt'。

7. Determine the new best metric / best model checkpoint

根据指定的指标来判断当前模型是否为最佳模型，并记录下最佳指标值和对应的模型检查点路径。

if metrics is not None and self.args.metric_for_best_model is not None:
    metric_to_check = self.args.metric_for_best_model
    if not metric_to_check.startswith("eval_"):
        metric_to_check = f"eval_{metric_to_check}"
    metric_value = metrics[metric_to_check]
    operator = np.greater if self.args.greater_is_better else np.less
    if (self.state.best_metric is None or self.state.best_model_checkpoint is None or operator(metric_value, self.state.best_metric)):
        self.state.best_metric = metric_value
        self.state.best_model_checkpoint = output_dir

8. self.state 状态保存

只要满足 self.args.should_save 条件，调用 self.state.save_to_json 来保存 state 的状态。

if self.args.should_save:
    self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))

9. 随机状态数保存

在非分布式训练中保存随机数生成器状态（Python、NumPy、PyTorch、CUDA、TPU）。根据进程数量决定保存为单个文件或多个文件。

rng_states = {
    "python": random.getstate(),
    "numpy": np.random.get_state(),
    "cpu": torch.random.get_rng_state(),
}
# ... cuda/xla handling ...
os.makedirs(output_dir, exist_ok=True)
if self.args.world_size <= 1:
    torch.save(rng_states, os.path.join(output_dir, "rng_state.pth"))
else:
    torch.save(rng_states, os.path.join(output_dir, f"rng_state_{self.args.process_index}.pth"))

10. 控制权重限制保存

循环查看保存权重数量限制功能。如果设置了 save_total_limit，当检查点数量超过限制时，会删除旧的检查点。

if self.args.should_save:
    self._rotate_checkpoints(use_mtime=True, output_dir=run_dir)

四、self.save_model(output_dir, _internal_call=True)

1. save_model 源码

def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False):
    if output_dir is None:
        output_dir = self.args.output_dir
    if is_torch_tpu_available():
        self._save_tpu(output_dir)
    elif is_sagemaker_mp_enabled():
        # Calling the state_dict needs to be done on the wrapped model and on all processes.
        ...
    elif (ShardedDDPOption.ZERO_DP_2 in self.args.sharded_ddp or ShardedDDPOption.ZERO_DP_3 in self.args.sharded_ddp or self.fsdp is not None or self.is_fsdp_enabled):
        ...
    elif self.is_deepspeed_enabled:
        ...
    elif self.args.should_save:
        self._save(output_dir)
    # Push to the Hub when `save_model` is called by the user.
    if self.args.push_to_hub and not _internal_call:
        self.push_to_hub(commit_message="Model save")

2. _save 源码解读

此部分保存 config.json 文件、模型权重文件 pytorch_model.bin 与 trainer 的 self.args 内容 training_args.bin。

def _save(self, output_dir: Optional[str] = None, state_dict=None):
    output_dir = output_dir if output_dir is not None else self.args.output_dir
    os.makedirs(output_dir, exist_ok=True)
    logger.info(f"Saving model checkpoint to {output_dir}")
    supported_classes = (PreTrainedModel,)
    if not is_peft_available():
        supported_classes = (PreTrainedModel, PeftModel)
    
    if not isinstance(self.model, supported_classes):
        if state_dict is None:
            state_dict = self.model.state_dict()
        if isinstance(unwrap_model(self.model), supported_classes):
            unwrap_model(self.model).save_pretrained(
                output_dir, state_dict=state_dict, safe_serialization=self.args.save_safetensors
            )
        else:
            logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
    else:
        self.model.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)
    
    # Good practice: save your training arguments together with the trained model
    torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))

此外，若存在 tokenizer 对象，也会一并保存。TRAINING_ARGS_NAME 即保存 trainer 中的 self.args 内容。

五、自定义参数保存示例

可以通过修改 trainer.args 添加自定义参数，并在保存模型时一同持久化。

# 定义 Trainer 对象
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
)
trainer.args.a = 20  # 添加自定义参数
# 开始训练
trainer.train()
trainer.save_model("/huggingface_demo/out_dirs")

通过继承 Trainer 或修改 save_model 等方法，可灵活实现自定义参数的保存需求。