[BUG] Deepspeed does not update the model when using "Qwen/Qwen2.5-3B" and is fine with ""Qwen/Qwen2.5-1.%B""

[MiladInk]

Describe the bug
I know this sounds very weird. However, when I use the deepspeed to optimize a "Qwen/Qwen2.5-3B" model, the model does not update at all. The same exact training code works with "Qwen/Qwen2.5-1.5B". Also checked and optimizing "meta-llama/Llama-3.2-3B" does not work. The parameters remain exactly the same. However, by just setting "torch_adam" to true, the issue goes away.

[loadams]

Hi @MiladInk - can you specify what you mean when you say you set torch_adam to true? And can you share your DeepSpeed version?

[loadams]

Hi @MiladInk - following up on the request above?

[loadams]

Hi @MiladInk - I'm not able to repro this, closing for now, if you see this, please let us know and we can investigate more with more information.

[MiladInk]

Sorry for the late reply. I mean the following argument in the config needs to be set to torch.Adam:

deepspeed_config = {
    "bf16": {"enabled": True},
    "zero_optimization": {"stage": 2, "overlap_comm": False},
    "train_batch_size": EPISODES_PER_ITERATION,
    "train_micro_batch_size_per_gpu": PER_DEVICE_BATCH_SIZE,
    "gradient_accumulation_steps": EPISODES_PER_ITERATION // PER_DEVICE_BATCH_SIZE,
    "gradient_clipping": 1.0,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": LEARNING_RATE,
            "betas": (0.9, 0.999),
            "eps": 1e-8,
            "weight_decay": 0.0,
            "torch_adam": True,
        },
    },
}

and my deepspeed version is 0.16.4

[tang-bd]

Hi, any update on this? I'm encountering the same issue.