DeepSpeed训练Dolly2.0学习率未改善及衰减学习率使用方法咨询

一、学习率未按预期变化的原因 #

你的DeepSpeed配置里使用了WarmupLR调度器，这个调度器仅实现学习率预热功能——在指定的warmup步数内从warmup_min_lr升到warmup_max_lr，之后学习率会一直保持在warmup_max_lr不再变化。而你直接训练时用的应该是带衰减逻辑的调度器（比如线性衰减、余弦衰减），这就是两者学习率变化曲线不一致的核心原因。

另外配置中调度器的参数全设为"auto"，如果你的训练脚本没有正确传递原生训练的衰减相关参数，DeepSpeed无法自动推断出衰减逻辑，只会维持预热后的固定学习率。

二、DeepSpeed中配置衰减学习率的方法 #

要实现衰减学习率，需要替换调度器类型为支持衰减的版本，DeepSpeed常用的有WarmupLinearLR（线性衰减）和WarmupCosineLR（余弦衰减），以下是配置示例：

修改后的DeepSpeed配置（以余弦衰减为例） #

{
    "fp16": {
      "enabled": false
    },
    "bf16": {
      "enabled": true
    },
    "optimizer": {
      "type": "AdamW",
      "params": {
        "lr": "auto",
        "betas": "auto",
        "eps": "auto",
        "weight_decay": "auto"
      }
    },
    "scheduler": {
      "type": "WarmupCosineLR",
      "params": {
        "warmup_min_lr": 0,
        "warmup_max_lr": "auto",
        "warmup_num_steps": 1000,
        "total_num_steps": 10000  // 替换为你的训练总步数
      }
    },
    "zero_optimization": {
      "stage": 3,
      "offload_optimizer": {
        "device": "cpu",
        "pin_memory": true
      },
      "offload_param": {
          "device": "cpu",
          "pin_memory": true
      },
      "overlap_comm": true,
      "contiguous_gradients": true,
      "sub_group_size": 1e9,
      "reduce_bucket_size": "auto",
      "stage3_prefetch_bucket_size": "auto",
      "stage3_param_persistence_threshold": "auto",
      "stage3_max_live_parameters": 1e9,
      "stage3_max_reuse_distance": 1e9,
      "stage3_gather_16bit_weights_on_model_save": true
    },
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto",
    "steps_per_print": 2000,
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "wall_clock_breakdown": false
  }

关键配置说明 #

• 调度器类型：选择WarmupLinearLR或WarmupCosineLR，对应线性衰减和余弦衰减逻辑
• warmup_num_steps：设置学习率预热的步数，和你直接训练时的参数保持一致
• total_num_steps：必须明确设置训练总步数（总步数 = 训练样本数 / 全局batch_size * 训练轮数），这是计算衰减的核心参数
• warmup_min_lr：预热起始学习率，通常设为0或一个极小值
• warmup_max_lr：预热结束后的峰值学习率，设为"auto"可继承optimizer中的lr设置

如果要使用线性衰减，只需将scheduler.type改为"WarmupLinearLR"即可，参数配置逻辑一致。

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