训练Hugging Face Wave2Vec模型时遇NameError：'PartialState'未定义

问题重现 #

按照Hugging Face音频分类指南训练Wave2Vec模型时，运行以下代码触发NameError，提示名称'PartialState'未定义：

training_args = TrainingArguments(
    output_dir="my_awesome_mind_model",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=3e-5,
    per_device_train_batch_size=32,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=32,
    num_train_epochs=10,
    warmup_ratio=0.1,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    # eval_dataset=encoded_minds["test"],
    tokenizer=feature_extractor,
    compute_metrics=compute_metrics,
)

trainer.train()

错误信息：

NameError Traceback (most recent call last)
in <cell line: 1>()
1 training_args = TrainingArguments(
2 output_dir="my_awesome_mind_model",
3 evaluation_strategy="epoch",
4 save_strategy="epoch",
5 learning_rate=3e-5,

4 frames
/usr/local/lib/python3.10/dist-packages/transformers/training_args.py in _setup_devices(self)
1629 self._n_gpu = 1
1630 else:
1631 self.distributed_state = PartialState(backend=self.ddp_backend)
1632 self._n_gpu = 1
1633 if not is_sagemaker_mp_enabled():

NameError: name 'PartialState' is not defined

解决方案 #

这个错误的核心原因是transformers库依赖的accelerate库未安装，或者两者版本不兼容——PartialState是accelerate库中负责分布式训练状态管理的类，新版本transformers会调用它。

解决方法如下：

• 安装或升级兼容版本库：
  在终端执行命令：

  pip install --upgrade transformers accelerate
  

  这会确保两个库处于兼容的最新版本，自动解决依赖问题。
• 临时修复（不推荐长期使用）：
  如果暂时不想升级库，可以在代码开头手动导入PartialState：

  from accelerate import PartialState
  

内容的提问来源于stack exchange，提问作者Nikolai Reverger