Stop and resume yet match an uninterrupted run — what goes in the resume file

Goal

Make the result of training 60 steps at once and the result of atomically saving a resume file (weights, optimizer state, step count, and batch random number state) at step 30 and resuming not differ by a single bit. Measure how far they diverge if you leave out the optimizer state or the data position, and what half-precision saving changes.

Why it matters

If you save only the weights and resume training, the result changes without any error. AdamW remembers two moments for each parameter, and if that memory is lost, the first few steps right after resuming move at completely different sizes. If you rewind the batch order in the middle of an epoch, you see the first part twice and never see the latter part. Such differences barely show up in the loss curve, so they are easy to miss. That is why MiniMind's lm_checkpoint keeps the weights and the resume file separately, and writes both to temporary files and then renames them. This lab confirms exactly what that design protects by comparing against the result done at once.

Steps

1. Write /root/mm/ckpt/run.py, train 60 steps at once (seed 7, batch 8, length 128, lr 3e-3), and leave /root/mm/ckpt/straight.pth and /root/mm/ckpt/straight_log.csv (step,train_loss).
2. Stop the same training at step 30 and save the resume file /root/mm/ckpt/resume.pth by writing to .tmp and then swapping with os.replace (model, optimizer, step, and gen_state).
3. Resume from the resume file and continue from step 31 to step 60, and leave /root/mm/ckpt/resumed.pth and /root/mm/ckpt/resumed_log.csv. They must equal step 1.
4. Resume with only the optimizer state thrown away and leave /root/mm/ckpt/noopt.pth.
5. Resume with only the batch random number state thrown away (restarting from the initial seed) and leave /root/mm/ckpt/noskip.pth.
6. Save straight.pth in half precision to /root/mm/ckpt/straight_fp16.pth, and write the two file sizes and the logit difference to /root/mm/ckpt/fp16.json.
7. Count the number of elements of the weights and of the optimizer state in the resume file and write them to /root/mm/ckpt/optim.json.
8. In /root/mm/ckpt/report.md, write the three sections ## 무엇을 저장하나, ## 빠뜨리면, and ## 반정밀도 (the Korean headings mean "What to save", "If you leave something out", and "Half precision"), and include the ratio from step 7 and the fp16 file size (bytes) from step 6.

Notes

• 60 steps take about 5 seconds on a node. If you put --stop-at, --resume, --no-optim, and --no-skip in one script, steps 3–5 are one line each.
• If you draw batches with torch.randint(..., generator=g), you can save and restore "the next batch to draw" with g.get_state() and g.set_state(). This one line does what MiniMind's SkipBatchSampler does.
• Read the resume file with torch.load(경로, weights_only=True) (the placeholder is the file path). If you just read a checkpoint someone else gave you, the Python code inside it may run.
• Common mistakes: not fixing the seed before creating the model so that the initialization on the resumed side differs (it gets overwritten soon so it looks fine, but the generator state goes out of line), and recounting the learning rate schedule from 0 based on the resumed step.
• Sources: trainer_utils.py — lm_checkpoint and SkipBatchSampler · PyTorch — Saving and Loading · torch.load weights_only · os.replace

60 steps at once — the reference

Write /root/mm/ckpt/run.py, train 60 steps at once with seed 7, batch 8, length 128, and lr 3e-3 (MiniMind cosine, based on 60 total steps), and leave /root/mm/ckpt/straight.pth (state_dict) and /root/mm/ckpt/straight_log.csv (step,train_loss, 1–60).

It is the same as the training loop in the earlier module. This time, write the loss to six decimal places — you compare digit by digit against the resumed log in step 3.

Save atomically at step 30

Stop the same training at step 30, write {"model": state_dict, "optimizer": optimizer.state_dict(), "step": 30, "gen_state": 배치 생성기.get_state()} (the Korean name in it stands for the batch generator) to /root/mm/ckpt/resume.pth.tmp, and then swap it to /root/mm/ckpt/resume.pth with os.replace. The learning rate is still computed based on 60 total steps.

If you write everything to a temporary file and then rename it, a whole file remains even if it dies during the save. This is what MiniMind's lm_checkpoint does with ckp_tmp → os.replace. If .tmp is left behind, you skipped the rename.

Resume to step 60 — it must equal the reference

In a new process, read resume.pth to restore the weights, optimizer state, and batch generator state, train from step 31 to step 60, and leave /root/mm/ckpt/resumed.pth and /root/mm/ckpt/resumed_log.csv (31–60).

Build the model and optimizer exactly as in step 1, overwrite them with load_state_dict, and restore the generator with g.set_state(…). The grader compares all tensors of the two checkpoints and checks whether the maximum difference is 1e-6 or less — if you run on the same node, it is 0.

If you throw away the optimizer state

From the resume file, restore only the weights and the batch generator state, continue steps 31–60 with a newly created optimizer, and leave /root/mm/ckpt/noopt.pth.

AdamW's step size is proportional to the gradient divided by the square root of its moving average. If the memory starts again from 0, the first few steps right after resuming move at completely different sizes. The grader checks whether this result is clearly different from the reference.

If you rewind the batch order

From the resume file, restore only the weights and the optimizer state, continue steps 31–60 with the batch generator rebuilt from the initial seed (7), and leave /root/mm/ckpt/noskip.pth.

If you leave the generator in its initial state, at step 31 you see the batch of step 1 again. This is why, when resuming in the middle of an epoch, MiniMind skips the batches already seen with SkipBatchSampler.

If you save in half precision

Convert all tensors of straight.pth with .half() and save to /root/mm/ckpt/straight_fp16.pth, and write the sizes of the two files (bytes) and the maximum absolute difference of the logits on the first 128 tokens of /opt/mm/ref/val.npy to /root/mm/ckpt/fp16.json as fp32_bytes, fp16_bytes, and max_logit_diff.

mmkit.load_model expands even fp16 weights to fp32 when loading them. The file shrinks by nearly half (it is not exactly half because it also contains things like names besides the tensors), but the logits differ at about the third decimal place.

Why the resume file is big

From resume.pth, write the number of elements of the weights (excluding lm_head.weight because it is shared with the embedding), the number of elements of the optimizer state tensors (those with one or more dimensions), and their ratio to /root/mm/ckpt/optim.json as model_numel, optimizer_numel, and ratio.

optimizer.state_dict()['state'] holds step (0-dimensional), exp_avg, and exp_avg_sq for each parameter. If the ratio is 2, the resume file is about three times the weights file — this is why checkpoints of large models are heavy.

Write down what to save and why

In /root/mm/ckpt/report.md, write the three sections ## 무엇을 저장하나, ## 빠뜨리면, and ## 반정밀도 (the Korean headings mean "What to save", "If you leave something out", and "Half precision"), and include the ratio from step 7 and the fp16_bytes from step 6 as numbers.

It is good to also write the divergence sizes seen in steps 4 and 5. In the half precision section, try writing one line on how to separate "weights for inference" from "weights for resuming".