[WIP][RFC] Always flatten model state_dict

[fegin]

The model state_dict is unique compared to other state dictionaries (e.g., optimizer). It's the only one that will be exported outside of TorchTitan and imported from other sources. To ensure FQN consistency, we previously removed the prefix during the first checkpoint load and last checkpoint save. However, this approach has caused confusion among users, despite available options to control behavior.

This PR aims to resolve the issue by always flattening the model state dictionary, eliminating the "MODEL." prefix from its keys. We decided not to flatten all components due to the risk of key collisions between different components. Instead, this PR only flattens the model state_dict, which is a special case.

While this solution isn't perfect, as it introduces different handling for different components, it's a good compromise given the unique nature of the model state_dict.

Also see the discussion in #1321 (comment)

This is the pseudo code for the current state:

if model_only:
    state_dict = model.state_dict()
else:
    state_dict = {
        "MODEL": model,
        "OPTIMIZER": optimizer,
         ...
     }
}

This is the pseudo code after this PR is landed:

state_dict = model.state_dict()
if not model_only:
    state_dict.update(
        {"OPTIMIZER": optimizer}
         ...
     )

FSDP4 v.s. FSDP4 TP2 loss curve with seed checkpoint and --training.seed=42

[tianyu-l]

could you help verify "save seed checkpoint -> load from seed checkpoint with different parallelism" still yields identical loss, as an additional integration test?

[tianyu-l]

does this mean "model" still exists in state_dict as a key -- we only flatten it in the checkpoint (and its load and save)

[fegin]

In Checkpointer, we still keep separate keys in self.states, like MODEL, OPTIMIZER. This will allow use to manipulate different state_dicts. This line use MODEL to access only the model state_dict but this line does not wrap the model state_dict, so there will be no model. prefix.

[fegin]

@tianyu-l The loss curves don't match, w/ or w/o freqs_cis in the seed checkpoint. The seed checkpoint may have broken before freqs_cis being removed from checkpoint.

[tianyu-l]

In terms of loss curve matching

local_batch_size is 2 when TP is 2, otherwise local_batch_size is 1.

In order to make sure dataloader behaves consistently across multiple runs, we need to fix the DP degree (dp_replicate * dpshard).
So a typical comparison you may have (by keeping the overall DP degree 4) could be

• FSDP 4
• DP 2, FSDP 2, TP 2
• DP 2, FSDP 2, CP 2, PP 2

For details and more examples, please see
https://github.com/pytorch/torchtitan/blob/main/docs/converging.md

@wwwjn we probably should include this in #1363

[fegin]

@tianyu-l I also tried FSDP4 v.s. FSDP4 TP2 before I tried this setting and the loss curves didn't match.

[fegin]

Correct the statement, FSDP4 v.s. FSDP4 TP2 match, with or without fixing `training.seed`.

[tianyu-l]

SGTM!