Save and load train_state.pt in checkpoints

Parent initiative: #2319

Depends on #3555 for the `TrainState` runtime contract.

## Context

`TrainState` will be the runtime source of truth for mutable training progress. This task covers persisting and restoring that state through Megatron-LM checkpoints while preserving compatibility with older checkpoints.

## Scope

- Save the active `TrainState` as `train_state.pt` in the checkpoint iteration directory.
- Load `train_state.pt` into the active `TrainState` when the sidecar is present.
- Make the loaded state available consistently across distributed ranks.
- Prefer `train_state.pt` when both the sidecar and legacy progress fields are present.
- When `train_state.pt` is absent, construct `TrainState` from the legacy checkpoint fields/`state_dict["args"]`.
- Verify that the sidecar and legacy fields contain equivalent values while both are written.
- Leave the distributed checkpoint representation for model, optimizer, scheduler, RNG, and rerun state unchanged.

## Completion criteria

- New Megatron-LM checkpoints contain `train_state.pt`.
- Synchronous and asynchronous checkpoint saves write the same training-state metadata.
- Saving and loading a checkpoint recovers every `TrainState` field.
- New Megatron-LM code reconstructs `TrainState` from older checkpoints without the sidecar.
- Tests cover round trips, legacy checkpoint format, sidecar precedence, and agreement between the two representations.

## Out of scope

- Defining the `TrainState` schema or lifecycle.
- Migrating training, evaluation, or logging code away from mutable fields in `args`.
- Saving or loading `PretrainConfigContainer`.

Start by tracing the checkpoint save and load paths and reading the TrainState runtime contract from #3555. Implement and test train_state.pt handling in the checkpoint iteration directory, including synchronous and asynchronous saves, distributed rank consistency, legacy reconstruction, precedence, and equivalent sidecar and legacy values.

Written by the indexing model from the issue text.