Pytorch Lightning 和 HuggingFace 的 Trainer 哪个好用？

import torch
from transformers import TrainerCallback

rank = 64
alpha = rank * 2
lora_config = LoraConfig(
    task_type = TaskType.CAUSAL_LM,
    target_modules = ["q_proj","k_proj","v_proj","o_proj", "gate_proj", "up_proj", "down_proj"],
    inference_mode = False, # 训练模式
    r=rank,
    lora_alpha=alpha, #Lora aLaph，具体作用参见Lora原理
    lora_dropout=0.1
)

train_args = TrainingArguments(
    output_dir=f"./output/Qwen/test_rank{rank}_alpha{alpha}",
    per_device_train_batch_size=1,
    gradient_accumulation_steps=3,
    logging_steps=4,
    num_train_epochs=7,    # 小模型要多训练会儿
    save_steps=100,
    eval_steps=10,
    learning_rate=8e-5,
    save_on_each_node=True,
    gradient_checkpointing=True,
    lr_scheduler_type="cosine",
    warmup_steps=20,
    report_to=["tensorboard"]
)

# Load best model at end=True
logger.info(f"args: {train_args}")
logger.info(f"【Start Training!】")

# 不LoRA
# model = get_peft_model(model, lora_config)
# model.print_trainable_parameters()

class MemoryTraceCallback(TrainerCallback):
    def __init__(self, output_dir="memory_trace", target_step=42):
        self.output_dir = output_dir
        self.target_step = target_step
        self.recorded = False

    def on_train_begin(self, args, state, control, **kwargs):
        # 开启显存历史记录 (包含堆栈信息)
        print(f"🚀 [Memory] 开始记录显存分配历史...")
        torch.cuda.memory._record_memory_history(
            max_entries=100000,
            context="all"  # 记录 Python 堆栈
        )

    def on_step_end(self, args, state, control, **kwargs):
        # 我们只需要跑几步，比如第 5 步结束后抓取快照
        if state.global_step == self.target_step and not self.recorded:
            print(f"📸 [Memory] 第 {self.target_step} 步结束，正在保存显存快照...")
            try:
                # 保存快照
                torch.cuda.memory._dump_snapshot(f"{self.output_dir}_snapshot.pickle")
                print(f"✅ [Memory] 快照已保存至 {self.output_dir}_snapshot.pickle")

                # 停止记录
		        torch.cuda.memory._record_memory_history(enabled=None)
                self.recorded = True
                # 可选：直接停止训练，因为我们只想要分析数据
                control.should_training_stop = True
            except Exception as e:
                print(f"❌ [Memory] 保存失败: {e}")


# 在 Trainer 中添加这个 Callback
trainer = Trainer(
    model=model,
    args=train_args,
    train_dataset=tokenized_id,
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),
    callbacks=[MemoryTraceCallback()],
)
trainer.train()

>>> features = [
...     {"input_ids": [1, 2, 3], "attention_mask": [1, 1, 1]},
...     {"input_ids": [4, 5], "attention_mask": [1, 1]},
... ]

>>> tokenizer.pad(features, return_tensors="pt")
{
  "input_ids": tensor([[1, 2, 3], [4, 5, 0]]),
  "attention_mask": tensor([[1, 1, 1], [1, 1, 0]])
}

class DataCollatorWithPadding:
    """Data collator that will dynamically pad the inputs received."""
    tokenizer: PreTrainedTokenizerBase
    padding: Union[bool, str, PaddingStrategy] = True
    max_length: Optional[int] = None
    pad_to_multiple_of: Optional[int] = None
    return_tensors: str = "pt"

    def __call__(self, features: list[dict[str, Any]]) -> dict[str, Any]:
        batch = pad_without_fast_tokenizer_warning(
            self.tokenizer,
            features,
            padding=self.padding,
            max_length=self.max_length,
	        pad_to_multiple_of=self.pad_to_multiple_of,
            return_tensors=self.return_tensors,
        )
        if "label" in batch:
            batch["labels"] = batch["label"]
            del batch["label"]
        if "label_ids" in batch:
            batch["labels"] = batch["label_ids"]
            del batch["label_ids"]

        return batch

if (self.label_smoother is not None or self.compute_loss_func is not None) and "labels" in inputs:
    labels = inputs.pop("labels")
else:
    labels = None

if self.model_accepts_loss_kwargs:
    kwargs = {}
    if num_items_in_batch is not None:
        kwargs["num_items_in_batch"] = num_items_in_batch
    inputs = {**inputs, **kwargs}
outputs = model(**inputs)

def ForCausalLMLoss(
    logits,
    labels,
    vocab_size: int,
    num_items_in_batch: Optional[torch.Tensor] = None,
    ignore_index: int = -100,
    shift_labels: Optional[torch.Tensor] = None,
    **kwargs,
) -> torch.Tensor:

    # Upcast to float if we need to compute the loss to avoid potential precision issues
    logits = logits.float()
    if shift_labels is None:
        # Shift so that tokens < n predict n
        labels = nn.functional.pad(labels, (0, 1), value=ignore_index)
        shift_labels = labels[..., 1:].contiguous()

    # Flatten the tokens
    logits = logits.view(-1, vocab_size)
    shift_labels = shift_labels.view(-1)

    # Enable model parallelism
    shift_labels = shift_labels.to(logits.device)
    loss = fixed_cross_entropy(logits, shift_labels, num_items_in_batch, ignore_index, **kwargs)

    return loss

def fixed_cross_entropy(
    source: torch.Tensor,
    target: torch.Tensor,
    num_items_in_batch: Optional[torch.Tensor] = None,
    ignore_index: int = -100,
    **kwargs,
) -> torch.Tensor:
    reduction = "sum" if num_items_in_batch is not None else "mean"
    loss = nn.functional.cross_entropy(source, target, ignore_index=ignore_index, reduction=reduction)
    if reduction == "sum":
        # just in case users pass an int for num_items_in_batch, which could be the case for custom trainer
        if torch.is_tensor(num_items_in_batch):
            num_items_in_batch = num_items_in_batch.to(loss.device)
        loss = loss / num_items_in_batch
    return loss