From ef1b586f47720d59a4cc7475da73ca9f7c434ca1 Mon Sep 17 00:00:00 2001 From: artem Date: Tue, 14 Jul 2026 15:30:05 +0000 Subject: [PATCH] Handle `step` key correctly in dp_reshardable checkpoint save with `--optimizer-cpu-offload` Graft of NVIDIA/Megatron-LM#2874 (f4502eb1c9): wrap the optimizer `step` in LocalNonpersistentObject in sharded_param_state_dp_reshardable. Co-Authored-By: Claude Opus 4.8 (1M context) --- megatron/core/optimizer/distrib_optimizer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 058fb05e54a..fa850b7c856 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -1705,6 +1705,9 @@ def sharded_param_state_dp_reshardable( tensors[key] = LocalNonpersistentObject(tensors[key]) continue if key == 'step': + # The optimizer state of STEP is a 0-dim tensor and is handled + # separately via param_groups, not as part of the gradient buffer. + tensors[key] = LocalNonpersistentObject(tensors[key]) continue assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), ( tensors[key].shape,