用 Huggingface 训练 Google Big Bird 的记忆力不足答案

【问题标题】：Running out of Memory Training Google Big Bird with Huggingface用 Huggingface 训练 Google Big Bird 的记忆力不足
【发布时间】：2022-01-23 23:51:13
【问题描述】：

由于内存不足错误，我一直在努力使用 Huggingface 转换器库来训练 Google 的 Big Bird 模型。我有两个 Tesla V100 GPU，每个都有 32 GB RAM。我正在尝试使用 Huggingface 训练器 API 在 Spider（SQL 数据集的自然语言）上训练 google/bigbird-roberta-base 模型 (https://huggingface.co/google/bigbird-roberta-base)。我正在使用 1 的批量大小和此模型的最小版本，但仍然出现 OOM 错误。根据 Big Bird 论文 (https://arxiv.org/abs/2007.14062)，Big Bird 可以在 16 GB 内存的芯片上进行训练，所以我不确定我为什么会遇到 OOM。有没有人因为记忆问题而在训练 Big Bird 时遇到困难？

这是进行训练的代码：

rouge = datasets.load_metric("rouge")

training_args = Seq2SeqTrainingArguments(
    predict_with_generate = True,
    evaluation_strategy = "steps",
    per_device_train_batch_size = batch_size,
    per_device_eval_batch_size = batch_size,
    output_dir = "./",
    logging_steps = 2,
    save_steps = 400,
    eval_steps = 4
)

def compute_metrics(pred):
    labels_ids = pred.label_ids
    pred_ids = pred.predictions

    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    labels_ids[labels_ids == -100] = tokenizer.pad_token_id
    label_str = tokenizer.batch_decode(labels_ids, skip_special_tokens=True)

    rouge_output = rouge.compute(predictions=pred_str, references=label_str, rouge_types=["rouge2"])["rouge2"].mid

    return {
        "rouge2_precision": round(rouge_output.precision, 4),
        "rouge2_recall": round(rouge_output.recall, 4),
        "rouge2_fmeasure": round(rouge_output.fmeasure, 4),
    }

trainer = Seq2SeqTrainer(
    model = model,
    tokenizer = tokenizer,
    args = training_args,
    compute_metrics = compute_metrics,
    train_dataset = train_data,
    eval_dataset = val_data
)

trainer.train()

这是我得到的确切错误：

RuntimeError: CUDA out of memory. Tried to allocate 36.00 MiB (GPU 0; 31.75 GiB total capacity; 25.14 GiB already allocated; 21.50 MiB free; 26.23 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

非常感谢您分享您在这方面的任何经验！

【问题讨论】：

标签： nlp web-crawler huggingface-transformers

【解决方案1】：

我对具有 48gb 内存的 a6000 也有同样的问题，在将批量大小从 128 oom 设置为 1 之后，我猜模型比他们说的要大。

【讨论】：

【解决方案2】：

我猜是抱脸的问题。 https://github.com/google-research/bigbird中的原始代码不存在OOM问题。

【讨论】：

正如目前所写，您的答案尚不清楚。请edit 添加其他详细信息，以帮助其他人了解这如何解决所提出的问题。你可以找到更多关于如何写好答案的信息in the help center。