Add model merging callback with integration test

[baileykuehl]

Added a model merging callback that averages model weights over the last N steps of training and saves the result as a stepXXXX-merged directory

Key features:

• Disabled by default (enabled=False) because it seems to somewhat affect performance at merge steps on larger models (measuring this impact was out of scope for now, will follow-up with experiments if this gets adopted)
• For WSD-S schedulers, automatically merges and evaluates right before each decay phase begins. This overlaps with the pre-decay eval in the old ladder.
• Saves accumulator state to checkpoints, so interrupted runs can resume mid-merge-window without losing progress
• Integration test added that trains a small model for 5 steps (averages the last 3 steps)

A few examples of recent runs with this callback:

• 60M: https://beaker.org/orgs/ai2/workspaces/oe-t-ladder/work/01KG625YKG9PB8Q753P327MB9R
• 370M: https://beaker.org/orgs/ai2/workspaces/oe-t-ladder/work/01KG626695QSZ6D30T41B3YJYQ

[epwalsh]

The way the trainer saves this state is not ideal for model state, especially when it's distributed/sharded. So we'll probably need to rethink this.

Here's one alternative:

Say we're merging the last 100 steps. We enforce saving a full training checkpoint right when start to accumulate, i.e. 100 steps before we want to evaluate the merged model, and we don't checkpoint again until after the merged model is evaluated. That way we wouldn't need to save the state of this call back because it would always be recomputed.

[baileykuehl]

That makes a lot of sense! So the trade-off of this would just be that if we get interrupted during the merge window, we would lose up to merge_last_n_steps - 1 steps of work, at most? Which seems reasonable while we're using fairly low values

[epwalsh]

Using dist_cp* functions comes with additional overhead that I think we can safely avoid by just collecting model parameters directly, e.g. something like model_state = {k: get_local_tensor(p.data.detach()) for k, p in model.named_parameters()}.

There's still the overhead of copying those params to CPU... One option is to keep them on GPU if we can afford it. But for larger models we're typically memory-starved, so that's doesn't scale. Alternatively we could still copy to CPU but try to hide the overhead of the copy by doing it asynchronously.

For example,

1. Trigger an async copy as early as possible, like in pre_load_batch(). You can do an async copy with olmo_core.utils.move_to_device(tensor, torch.device("cpu"), non_blocking=True). Now you can't do anything else with the CPU-tensor until you've triggered a host-device sync, so for now you wait.
2. Right before the optimizer step you force a host-device sync (call torch.cuda.synchronize()) to ensure the copy to CPU has finished before params are modified.
3. Then accumulate those weights.

This means the accumulated weights are a step behind (since they're updated pre-optim step), but I think we can adjust the bookkeeping accordingly.