[Bug] Checkpoint resume rejects valid partially initialized AdamW state

[Ideny42]

System info

source /data/unirl-dyref/venv-sglang/bin/activate

git rev-parse HEAD
python -c "import sys, platform; print(sys.version); print(platform.platform())"
python -c "import torch; print('torch', torch.version, 'cuda', torch.version.cuda); print(torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'no cuda')"
python -c "import importlib.metadata as m
for pkg in ('unirl', 'vllm', 'vllm-omni', 'sglang', 'torch'):
try:
print(pkg, m.version(pkg))
except m.PackageNotFoundError:
print(pkg, 'not installed')"

Rollout / inference engine

Train-side only

Domain

sft (train_sft)

Recipe / config-name

N/A (generic checkpoint reproducer; first observed in a modified Qwen-Image-Edit SFT recipe)

How you ran it

• Official example in examples/ (unmodified recipe)
• Modified recipe or my own script

Reproduction

Summary

UniRL's optimizer checkpoint save/load path is not round-trip safe when an
AdamW optimizer is partially initialized:

• some parameters have received gradients and own AdamW state;
• other parameters are in optimizer param_groups but have never received a
  gradient and legitimately have no state entry.

The checkpoint can be saved successfully, but restoring it through
set_optimizer_state_dict fails because the loader requires a state entry for
every parameter listed in param_groups.

This is a generic checkpoint-boundary issue. It was first exposed by a
not-yet-merged Qwen-Image-Edit SFT recipe, but it is not caused by that recipe
or by SFT.

Minimal CPU reproduction

This reproducer does not depend on Qwen, LoRA, FSDP, DyRef, a dataset or a GPU:

import torch
from torch import nn
from torch.distributed.checkpoint.state_dict import (
    StateDictOptions,
    get_optimizer_state_dict,
    set_optimizer_state_dict,
)


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.used = nn.Linear(2, 2)
        self.unused = nn.Linear(2, 2)

    def forward(self, x):
        return self.used(x)


model = Model()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)

model(torch.ones(1, 2)).sum().backward()
optimizer.step()

options = StateDictOptions(full_state_dict=True, cpu_offload=True)
state_dict = dict(
    get_optimizer_state_dict(model, optimizer, options=options)
)

params = [
    name
    for group in state_dict["param_groups"]
    for name in group["params"]
]
missing = [
    name
    for name in params
    if name not in state_dict["state"]
]

print(
    "state", len(state_dict["state"]),
    "params", len(params),
    "missing", missing,
)

restored_model = Model()
restored_optimizer = torch.optim.AdamW(
    restored_model.parameters(),
    lr=1e-3,
)

set_optimizer_state_dict(
    restored_model,
    restored_optimizer,
    optim_state_dict=state_dict,
    options=StateDictOptions(full_state_dict=True),
)

Observed with Python 3.11.6 and torch 2.5.0:

state 2 params 4 missing ['unused.weight', 'unused.bias']
KeyError: 'unused.weight'

unused.weight and unused.bias are valid AdamW parameters. They simply have
not received a gradient yet, so AdamW has not lazily created step, exp_avg
or exp_avg_sq for them.

UniRL code path

UniRL uses this same pair of APIs:

• gather_optimizer_state_dict /
  sharded_optimizer_state_dict call get_optimizer_state_dict;
• load_optimizer_state_dict /
  load_sharded_optimizer_state_dict call set_optimizer_state_dict;
• BaseFSDP2Backend saves the resulting optimizer state and restores it
  unconditionally during training-state resume.

As a result, UniRL can emit an optimizer checkpoint that its own resume path
cannot consume.

Static audit was performed against:

origin/main c8af5d699c81fb02a0086e34b342b1efacc7d431

Relation to PR #440

PR #440 fixed a related but narrower case: exporting an entirely cold AdamW
optimizer without advancing its first-step clock.

Its guard requires:

not optimizer.state

That condition is false for a partially initialized optimizer: once any
parameter has been updated, optimizer.state is non-empty even if other
parameters still have no state.

Therefore the partially initialized case passes through the #440 export
boundary unchanged and remains unhandled.

Related PR:
#440

Qwen-Image trigger

The issue was first observed while resuming a Qwen-Image-Edit LoRA SFT run.

The final QwenImage transformer block still consumes text K/V for image
attention, but its updated text output is discarded after the block. LoRA
parameters attached only to the final text-output path consequently keep
grad=None.

The observed optimizer contained:

1440 parameter names in param_groups
1436 entries in state
4 valid, never-initialized entries absent from state

The missing entries were LoRA A/B under these final-block modules:

transformer_blocks.59.attn.to_add_out
transformer_blocks.59.txt_mlp.net.2

Existing merged Qwen-Image and Qwen-Image-Edit training recipes already target
the same terminal attn.to_add_out module, for example:

examples/diffusion/qwen_image/qwen_image_trainside.yaml
examples/diffusion/qwen_image_edit_plus/qwen_image_edit_plus_flowgrpo_sglang.yaml

Their resume paths have not yet been runtime-verified for this failure, so they
are potential affected paths rather than confirmed reproductions.

The custom SFT recipe made the partially initialized state deterministic and
exposed the existing checkpoint assumption; it did not introduce that
assumption.

Workaround validation

Adding only an empty dictionary for each missing optimizer-state entry makes
the same checkpoint load successfully:

for group in state_dict["param_groups"]:
    for name in group["params"]:
        state_dict["state"].setdefault(name, {})

Validation results:

load: OK
future lazy init: ['exp_avg', 'exp_avg_sq', 'step']

On the real Qwen run:

• only 4 / 1440 empty entries were added;
• every existing optimizer tensor was preserved unchanged;
• model, scheduler and data-cursor state were unchanged;
• resume restored step 2517 and data epoch 5.000;
• the first resumed update had finite loss and gradients;
• the subsequently saved checkpoint at step 3021 resumed successfully.

This preserves standard AdamW lazy-initialization semantics.

Impact

Any UniRL training path can encounter this when its optimizer contains a
parameter that does not receive a gradient before checkpointing. Possible
causes include:

• an intentionally unused terminal branch;
• conditional routing;
• sparse experts;
• optional heads;
• parameters activated only by certain batches;
• broad LoRA target matching.

The failure is delayed until resume, so training and checkpoint saving both
appear successful.

This does not affect weight-only or exported-adapter inference loading.

The torch-format path is confirmed. The DCP-format path should be audited but
is not yet claimed as runtime-confirmed.

Proposed direction

Normalize valid sparse optimizer state at the shared checkpoint boundary:

• use optimizer param_groups as the authoritative parameter inventory;
• preserve every existing optimizer-state entry exactly;
• materialize {} only when the entire state entry for a listed parameter is
  absent;
• do not synthesize step, exp_avg or exp_avg_sq;
• continue rejecting malformed partially present state;
• preserve PR fix(checkpoint): preserve cold AdamW state during checkpoint export #440's fully cold AdamW step-zero semantics;
• support loading already-created sparse checkpoints.

This should be a generic optimizer checkpoint fix, not a Qwen-, LoRA-,
algorithm- or recipe-specific exception.

Acceptance criteria

• A save → resume cycle succeeds when one AdamW parameter has updated state and
  another has never received a gradient.
• Existing step, exp_avg and exp_avg_sq tensors are unchanged.
• A previously unused parameter still initializes normal AdamW state on its
  first future gradient.
• The fully cold AdamW behavior covered by fix(checkpoint): preserve cold AdamW state during checkpoint export #440 remains unchanged.
• Strict model-state validation remains unchanged.
• Existing sparse checkpoints remain loadable.
• Torch and DCP checkpoint paths are both audited; any unverified path is
  documented explicitly.

Your contribution

I plan to work on a focused root fix and submit a PR after maintainer
confirmation.

I will keep the change in the shared optimizer checkpoint boundary rather than
special-casing Qwen-Image or modifying individual recipes. Verification
commands and results will be included in the PR Test Plan; no one-off test
directory will be committed.

AI assistance: the investigation and proposed fix were developed with AI
assistance. I reviewed the checkpoint contents, the origin/main checkpoint
path, the relationship to #440, and the workaround behavior.

Expected behavior

A checkpoint emitted by UniRL must be consumable by UniRL's corresponding
resume path.

AdamW parameters listed in optimizer param_groups may legitimately have no
state entry when they have never received a gradient. Resume should preserve
such parameters as lazily uninitialized while restoring every existing
optimizer tensor, scheduler state, optimizer-step count, model state and
trainer state unchanged.

The fix should apply at the generic optimizer checkpoint boundary, preserve the
fully cold behavior from #440, remain backward-compatible with existing sparse
checkpoints, and not weaken strict model-state validation.