Loss metrics dramatically change after resuming from checkpoint

[darkmirage]

I am consistently getting bad loss values when resuming from a checkpoint.

Green is the original run, red and blue resumes from different checkpoints.

On 16 GPUs.

On 32 GPUs.

Runs were done on the 8b config with these overriding flags:

--metrics.enable_wandb --training.compile --training.seq_len 8192 \
--training.batch_size 1 --training.tensor_parallel_degree 1 \
--training.data_parallel_shard_degree -1 --training.data_parallel_replicate_degree 1 \
--training.steps 5000 --training.warmup_steps 200 --activation_checkpoint.mode selective \
--checkpoint.interval 500 --training.seed 42

[darkmirage]

It case it helps to narrow down the issue, it does appear that the shape of the curve of the resumed run is influenced by the config training.steps. It's not exactly the case, but the convergence point seems to be close to the total number of steps.

If I force the resumed run to step beyond training.steps the loss curve goes up very steeply.

[tianyu-l]

@darkmirage
Thanks for reporting the issue.

Can you be more specific on what were the configs before and after, if you ever changed them? E.g. in the post I only see one set of config, without checkpointing enabled. Please include a full config change.

Also I didn't understand

It's not exactly the case, but the convergence point seems to be close to the total number of steps.

In current setup, we expect resuming from exact same training.steps and warmup_steps, because the behavior of LR scheduler would depend on it. The behavior is undefined if you modify the steps config.

[darkmirage]

The config used with the default llama3_8b.toml file with those override flags added to the CLI.

In current setup, we expect resuming from exact same training.steps and warmup_steps, because the behavior of LR scheduler would depend on it. The behavior is undefined if you modify the steps config.

This is good to know because it seems a bit unintuitive as it would be common to want to train a checkpoint from a previous job further. What would be the suggested workflow for that? Creating a new checkpoint and reseting the LR scheduler config manually?

That said, the first two 16 GPU runs above were restarted from the checkpoints with the same steps and warmup_steps settings. I just meant that for cases where we do change the steps parameter because we want to train the model further, the shape of the mismatch curve seems to be influenced by this change.

Overall it sounds like something to do with the LR scheduler is the issue here.