Skip to content

Honor adam_w_mode in the CPU multi_tensor_adam binding #8307

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tohtana merged 5 commits into from
Aug 27, 2026
Merged

Honor adam_w_mode in the CPU multi_tensor_adam binding #8307

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
1e17e84
Honor adam_w_mode in the CPU multi_tensor_adam binding
PKUWZP Aug 24, 2026
9d209cc
Merge branch 'master' into cpu-adam-mode-fix
delock Aug 24, 2026
53949b5
Merge branch 'master' into cpu-adam-mode-fix
PKUWZP Aug 27, 2026
a4c407e
Address review: scope fp16 skip to TestAdamConfigs, name the placehol…
PKUWZP Aug 27, 2026
15dead9
Merge remote-tracking branch 'origin/cpu-adam-mode-fix' into cpu-adam…
PKUWZP Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 17 additions & 5 deletions csrc/cpu/adam/fused_adam.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,13 +19,25 @@ void multi_tensor_adam(int chunk_size,
const int bias_correction,
const float weight_decay)
{
static bool initialized = false;
if (!initialized) {
create_adam_optimizer(0);
initialized = true;
// ds_adam_step reads lr/betas/eps/weight_decay per call; only AdamW-vs-L2 is fixed at
// construction, so keep one optimizer instance per mode (mode 1 == AdamW, as in CUDA).
// The constructor hyperparameters are placeholders: update_state overwrites them on every
// step, so any value works here.
constexpr float kPlaceholderHyperparam = 0.0f;
static bool initialized[2] = {false, false};
Comment thread
PKUWZP marked this conversation as resolved.
const int optimizer_id = mode;
if (!initialized[mode]) {
create_adam_optimizer(optimizer_id,
kPlaceholderHyperparam,
kPlaceholderHyperparam,
kPlaceholderHyperparam,
kPlaceholderHyperparam,
kPlaceholderHyperparam,
mode == 1);
initialized[mode] = true;
}
for (int i = 0; i < tensor_lists[0].size(); i++) {
ds_adam_step(0,
ds_adam_step(optimizer_id,
step,
lr,
beta1,
Expand Down
6 changes: 4 additions & 2 deletions tests/unit/ops/adam/test_adamw.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,6 @@
from unit.simple_model import SimpleModel
from deepspeed.accelerator import get_accelerator

if torch.half not in get_accelerator().supported_dtypes():
pytest.skip(f"fp16 not supported, valid dtype: {get_accelerator().supported_dtypes()}", allow_module_level=True)
# yapf: disable
#'optimizer, zero_offload, torch_adam, adam_w_mode, resulting_optimizer
adam_configs = [["AdamW", False, False, False, (FusedAdam, True)],
Expand All @@ -38,6 +36,10 @@
@pytest.mark.parametrize(
'optimizer, zero_offload, torch_adam, adam_w_mode, resulting_optimizer',
adam_configs)
# Skipping at module level would also hide the dtype-parametrized reference test below, which is
# what let the CPU adam_w_mode bug go unnoticed on runners without fp16 support.
@pytest.mark.skipif(torch.half not in get_accelerator().supported_dtypes(),
reason=f"fp16 not supported, valid dtype: {get_accelerator().supported_dtypes()}")
class TestAdamConfigs(DistributedTest):
world_size = 1
reuse_dist_env = True
Expand Down
Loading