Optimizer state does not survive the stacked/unrolled layout change #21

Migrated from https://github.com/google-pytorch/torchtitan-archive/issues/389 (opened by jack8558 on 2026-09-11).

Issue mentioned from https://github.com/google-pytorch/torchtitan-archive/pull/356 and https://github.com/google-pytorch/torchtitan-archive/pull/355 comment (item 6).

When we stack layers, the parameter names change: layers.0.attention.wq.weight becomes layer_stacks.attention__wq__weight.

The model state dict is not affected, because state_dict runs inside unrolled(), which puts the old per-layer names back.

The optimizer state is affected, because it gets its names from named_parameters(), and that is not wrapped by unrolled(). Because of that, a stacked run and an unrolled run have different names for the optimizer state, causing a conflict when trying to load from a checkpoint.

It also fails silently, because _unflatten_optim_state_dict treats a missing name as "this parameter has no state" instead of an error. So if you save with --tpu-config.use-stacked-layers off and resume with it on, the weights come back fine but Adam starts from zero, with no warning.

Test to un-skip: test_layer_stack_reshard.py::test_the_optimizer_state_is_keyed_the_same_way_in_both_layouts.