使用DeepSpeed加速训练时调用model.eval()返回NoneType对象的问题求助

我最近在尝试用DeepSpeed加速模型训练，但在验证阶段调用model.eval()时遇到了一个诡异的问题：调用之后self.model居然变成了None！

问题相关代码片段 #

def evaluate(self, epoch_num=None, keep_all=True):
        print("self.model:", self.model)

        self.model = self.model.eval()
        print("self.model after eval:", self.model)

控制台输出日志 #

self.model: DeepSpeedEngine(
  (module): TSTransformerEncoder(
    (project_inp): Linear(in_features=6, out_features=128, bias=True)
    (pos_enc): LearnablePositionalEncoding(
      (dropout): Dropout(p=0.1, inplace=False)
    )
    (transformer_encoder): TransformerEncoder(
      (layers): ModuleList(
        (0-2): 3 x TransformerBatchNormEncoderLayer(
          (self_attn): MultiheadAttention(
            (out_proj): NonDynamicallyQuantizableLinear(in_features=128, out_features=128, bias=True)
          )
          (linear1): Linear(in_features=128, out_features=256, bias=True)
          (dropout): Dropout(p=0.1, inplace=False)
          (linear2): Linear(in_features=256, out_features=128, bias=True)
          (norm1): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (norm2): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (dropout1): Dropout(p=0.1, inplace=False)
          (dropout2): Dropout(p=0.1, inplace=False)
        )
      )
    )
    (output_layer): Linear(in_features=128, out_features=6, bias=True)
    (dropout1): Dropout(p=0.1, inplace=False)
  )
)
self.model after eval: None

没有用DeepSpeed的时候，模型的训练和评估都是正常的，一旦用上DeepSpeed就出现了这个问题。

我的DeepSpeed初始化方式 #

model, optimizer, _, _ = deepspeed.initialize(
        model=model,
        optimizer=optimizer,
        config_params=ds_config
    )

DeepSpeed配置文件（ds_config） #

{
    "fp16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
 
    "optimizer": {
        "params": {
            "lr": 0.001,
            "weight_decay": 0,
            "optimizer_class": "optimizers.RAdam"
        }
    },
 
    "zero_optimization": {
        "stage": 1,
        "overlap_comm": true,
        "contiguous_gradients": true
    },


    "zero_allow_untested_optimizer": true,
    "train_batch_size": 256,
    "steps_per_print": 2000,
    "wall_clock_breakdown": false
}

问题分析 #

我本来以为self.model.eval()只是把模型切换到评估模式，模型本身不会变成None，但实际日志显示调用后self.model就成了None。我怀疑这和DeepSpeed的封装或者配置有关，但具体是什么原因实在搞不清楚，有没有大佬能帮忙看看？

相关环境信息 #

• Python版本: 3.8.20
• PyTorch版本: 2.4.1
• DeepSpeed版本: 0.16.4

备注：内容来源于stack exchange，提问作者external