使用Google Colab调用HF PushToHubCallback训练崩溃问题求助

我正在使用Google Colab将微调后的模型推送到Hugging Face Hub。运行model.fit()函数时，Colab云端会创建输出目录，Bert-base-uncased模型开始在glue-sst2数据集上训练2个epoch。但当第二个epoch进度条完成后，训练仍持续数小时最终崩溃，当天Colab显示无法连接GPU后端。已尝试两次均出现相同问题，是否是Colab GPU资源不足或操作有误？

from datasets import load_dataset
raw_datasets = load_dataset('glue', 'sst2')

checkpoint = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = TFAutoModelForSequenceClassification.from_pretrained(checkpoint)

def tokenizer_function(example):
  return tokenizer(example['sentence'], truncation=True)

tokenized_datasets = raw_datasets.map(tokenizer_function, batched= True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer, return_tensors="tf")

tf_train_dataset = tokenized_datasets["train"].to_tf_dataset(
    columns=["attention_mask", "input_ids","token_type_ids"],
    label_cols=['labels'],
    shuffle=True,
    collate_fn=data_collator,
    batch_size=30,
)

tf_validation_dataset = tokenized_datasets["validation"].to_tf_dataset(
    columns=["attention_mask", "input_ids","token_type_ids"],
    label_cols=['labels'],
    shuffle=False,
    collate_fn=data_collator,
    batch_size=30,
)

num_epochs = 2
num_train_steps = len(tf_train_dataset) * num_epochs
lr_scheduler = PolynomialDecay(
    initial_learning_rate=2e-5, end_learning_rate=0.0, decay_steps=num_train_steps
)
opt = Adam(learning_rate=lr_scheduler)
loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(loss=loss, optimizer=opt, metrics= ['accuracy'])

callbacks = PushToHubCallback(
    "bert-fine-tuned_sst2", save_strategy="epoch", tokenizer=tokenizer
)

model.fit(tf_train_dataset, validation_data=tf_validation_dataset, epochs=1, callbacks=callbacks)  #this is where am having issues

• 修正代码逻辑矛盾：你定义了num_epochs=2，但model.fit()里设置的是epochs=1，参数不一致可能导致训练流程异常，先将model.fit()的epochs改为2，保持参数统一。
• 缓解Colab GPU资源压力：免费Colab的GPU有使用时长和显存限制，BERT-base搭配batch_size=30可能在训练后期（尤其是回调保存模型阶段）占用过量显存，引发崩溃。可以尝试：
  • 降低batch_size至16或8，减少单步训练的显存占用；
  • 在模型初始化前添加显存动态分配代码：import tensorflow as tf; tf.config.experimental.set_memory_growth(tf.config.list_physical_devices('GPU')[0], True)，避免显存被一次性占满。
• 拆分模型训练与推送流程：PushToHubCallback在epoch结束时推送模型的过程可能因网络或资源问题卡住，导致训练停滞。可以先移除回调，完成训练后再手动推送：

  # 训练完成后手动推送
  model.push_to_hub("bert-fine-tuned_sst2")
  tokenizer.push_to_hub("bert-fine-tuned_sst2")
  

• 提升Colab会话稳定性：训练期间不要关闭浏览器标签，可添加定期打印日志的代码（比如在训练循环中添加print语句），避免会话因闲置被系统回收。

内容的提问来源于stack exchange，提问作者Azeez Liadi