[BUG] Megatron Engine fails to save optimizer checkpoint

[daihaowz]

Checklist

• The error occurs when using our provided Docker image.
• I can consistently reproduce the bug across multiple trials or random seeds.
• If the error causes experiment abortion, I've verified that this error is the root
  cause, not a secondary error caused by peer workers.

Detailed Information

Describe the bug

When MegatronEngine saves a recovery checkpoint with optimizer state enabled (default behavior), MegatronCheckpointManager.generate_state_dict() calls self.optimizer.sharded_state_dict(state_dict) without specifying metadata, which defaults to the fully_sharded_model_space sharding strategy.

This strategy sets flattened_range=item_slice in distrib_optimizer.py:1760, but megatron-core 0.16.0+ added a blanket rejection of non-None flattened_range in ShardedTensor.validate_metadata_integrity() (mapping.py:134), causing an exception during optimizer state serialization.

This is an internal inconsistency in megatron-core 0.16.0+: the default sharding strategy uses a feature that the same version explicitly forbids.

Root cause location: areal/engine/megatron_utils/checkpointer.py:283

optimizer_sharded_states = self.optimizer.sharded_state_dict(state_dict)

Introduced by: bf9b3c3be (2026-03-17) bumped megatron-core from 0.13.1 to 0.16.0 without adapting the checkpointer. A workaround (no_save_optim flag) was added in 58304de84 (2026-03-23), but the root cause was never fixed.

Expected behavior

Recovery checkpoints should save optimizer state (Adam momentum, variance) successfully, allowing training to resume without losing optimizer state and requiring a warmup restart.

Full logs

Traceback (most recent call last):
  File ".../megatron/core/optimizer/distrib_optimizer.py", line 1766, in _get_param_state_sharded_tensors
    tensors[state_key].validate_metadata_integrity()
  File ".../megatron/core/dist_checkpointing/mapping.py", line 134, in validate_metadata_integrity
    raise CheckpointingException("ShardedTensor.flattened_range is not supported.")
megatron.core.dist_checkpointing.core.CheckpointingException: ShardedTensor.flattened_range is not supported.

To Reproduce

Commit ID

main branch

Environment

• megatron-core >= 0.16.0 (pyproject.toml currently declares 0.17.0)
• MegatronEngine with distributed optimizer enabled
• RecoverConfig with default no_save_optim=False

Script

Any RL training script using MegatronEngine that triggers a recovery checkpoint save. For example:

recover:
  mode: auto
  no_save_optim: false  # default value, triggers the bug

Fix: Pass metadata={"distrib_optim_sharding_type": "dp_reshardable"} in checkpointer.py:283. The dp_reshardable strategy sets flattened_range=None and works correctly under 0.16.0+. The only limitation is that DP configuration cannot change on resume, which is not a concern for recovery checkpoints.

[koladefaj]

This root cause was fixed in #1468 (89e05acb, 2026-07-02) — the analysis
here is accurate, and the proposed fix is exactly what landed. This issue
looks stale rather than open.

areal/engine/megatron_utils/checkpointer.py now reads:

optimizer_sharded_states = self.optimizer.sharded_state_dict(
    state_dict,
    is_loading=is_loading,
    metadata={"distrib_optim_sharding_type": "dp_reshardable"},
)

I verified the mechanism end to end against an installed megatron-core 0.17.0,
driving a real DistributedOptimizer on a single GPU the way
generate_state_dict() does:

megatron-core: 0.17.0
=== no metadata (the reported call) ===
FAIL -> megatron.core.dist_checkpointing.core.CheckpointingException:
        ShardedTensor.flattened_range is not supported.
=== metadata={'distrib_optim_sharding_type': 'dp_reshardable'} ===
OK
=== dp_reshardable, is_loading=True (load side) ===
OK

Every line-level claim in the report checks out on 0.17.0:
validate_metadata_integrity() ends in an unconditional
if self.flattened_range is not None: raise; sharded_state_dict defaults
sharding_type to 'fully_sharded_model_space' when no metadata is passed
(distrib_optimizer.py:1258); that strategy routes to
sharded_param_state_fs_model_space (:1701) → flattened_range=item_slice
(:1760) → raise (:1766); and sharded_param_state_dp_reshardable (:1551) sets
flattened_range=None (:1695), which is why the fix works.

Worth noting for anyone landing here from the traceback: #1468 also passes
is_loading=True when building the load-side template. Without it a freshly
built optimizer skips megatron-core's state pre-allocation
(distrib_optimizer.py:1276), the template only requests param, and DCP
silently drops exp_avg/exp_avg_sq on resume — training continues at full
learning rate with reset moments. So the save-side metadata alone is not
sufficient for a correct recovery.

One leftover: RecoverConfig.no_save_optim still documented itself as
"Required when using use_distributed_optimizer with Megatron (flattened_range
incompatibility)", which is the obsolete workaround and now silently costs
users their Adam moments across recovery. #1609 rewords that (and
no_load_optim) and adds regression tests.

Suggest closing this once #1609 lands, or now — the bug itself is already
fixed.