finetuning/sft_12hz.py: checkpoint save is non-atomic (can silently leave base weights) and peaks at 2x model size in host RAM

Description

Two independent defects in the checkpoint-saving block of finetuning/sft_12hz.py (lines 126-158 at 022e286). They share one small fix, so I'm reporting them together as one code path.

1. The save is not atomic, and the failure mode is a silent one.

output_dir = os.path.join(args.output_model_path, f"checkpoint-epoch-{epoch}")
shutil.copytree(MODEL_PATH, output_dir, dirs_exist_ok=True)   # :128  copies base model.safetensors IN
...
save_path = os.path.join(output_dir, "model.safetensors")
save_file(state_dict, save_path)                              # :158  overwrites it with trained weights

copytree writes the base model.safetensors into the final checkpoint directory, and save_file overwrites it ~30 lines later. Between those two statements, checkpoint-epoch-N/ is a complete, loadable checkpoint directory — correct config.json with tts_model_type: custom_voice and spk_id already patched in, all tokenizer files present — that contains untrained base weights.

If the process dies in that window (OOM, preemption, Ctrl-C, full disk), you are left with a checkpoint that loads without error, generates audio without error, and has learned nothing. There is no truncated file and no traceback to indicate it. On a spot/preemptible or Colab runtime this window is hit routinely.

This is what makes it worth reporting over the ordinary "crash mid-write" case: the artifact is not corrupt, it is wrong, and it looks fine.

2. Host RAM peaks at 2x the model size, because state_dict is never released.

state_dict = {k: v.detach().to("cpu") for k, v in unwrapped_model.state_dict().items()}   # :148

state_dict is a plain local inside the epoch loop. It is never deleted, so it stays bound through the following epoch's training and into the next save. On that next save, Python fully evaluates the dict comprehension on the right-hand side before rebinding the name — so a second complete CPU copy of every parameter exists simultaneously with the first.

Peak host RAM at every save after the first is therefore one full model in training dtype plus another one, not one.

The consequence is a crash that points nowhere near the cause: training runs fine, epoch 0 saves fine, and the process is killed by the OOM killer during the second save. We lost two runs to this before instrumenting ru_maxrss inside the save block. The signature to look for is "died at its second checkpoint, never a first".

3. Minor, same cause as 1: the copytree copies the base weight files in full and then discards them. For Qwen3-TTS-12Hz-0.6B-Base the base snapshot is 2.34 GiB across two weight files, and the one that is about to be overwritten is copied for nothing on every save. Proportionally more for the 1.7B.

Reproduction

The RAM behaviour is plain Python rebinding semantics, reproducible without the model:

live = set()
class Chunk:
    def __init__(self, tag): self.tag = tag; live.add(tag)
    def __del__(self): live.discard(self.tag)

sd = {i: Chunk(f"epoch0_{i}") for i in range(3)}    # first save

def build(tag):
    d = {i: Chunk(f"{tag}_{i}") for i in range(3)}  # RHS built in full...
    print("alive during build:", sorted(live))      # ...while the old dict is still bound
    return d

sd = build("epoch1")                                # ...only then rebound
# alive during build: ['epoch0_0', 'epoch0_1', 'epoch0_2', 'epoch1_0', 'epoch1_1', 'epoch1_2']

For the non-atomic save, the window is visible directly:

# during epoch 0's save, between :128 and :158
ls -la output/checkpoint-epoch-0/          # complete checkpoint dir
python -c "
from safetensors import safe_open
import hashlib
# model.safetensors here is byte-identical to the base model's until :158 runs
"

Interrupting anywhere in that window (or letting the OOM in defect 2 land there) leaves the directory in place, and a later Qwen3TTSModel.from_pretrained('output/checkpoint-epoch-0') succeeds with base weights.

Suggested fix

Both are addressed by writing to a temporary directory, renaming it into place only once the weights are written, skipping the base weight file in the copy, and freeing the state dict:

final_dir  = os.path.join(args.output_model_path, f"checkpoint-epoch-{epoch}")
output_dir = final_dir + ".tmp"
shutil.rmtree(final_dir,  ignore_errors=True)   # rename onto a non-empty dir is ENOTEMPTY
shutil.rmtree(output_dir, ignore_errors=True)

# don't copy the base weights just to overwrite them
_skip = lambda d, names: (["model.safetensors"]
    if os.path.abspath(d) == os.path.abspath(MODEL_PATH) else [])
shutil.copytree(MODEL_PATH, output_dir, dirs_exist_ok=True, ignore=_skip)

# ... existing config.json rewrite and save_file(state_dict, ...) ...

del state_dict, weight
gc.collect()
os.rename(output_dir, final_dir)                # atomic: the name appears only when complete

weight at :155 is also a plain local aliasing a tensor inside state_dict, so it needs to be released too.

I'm happy to open a PR for this if it would be useful.

Logs

Measured with resource.getrusage(0).ru_maxrss printed inside the save block, Qwen3-TTS-12Hz-0.6B-Base, weights in fp32:

before (upstream rebind pattern retained):
  1st save   OK          host RSS 9.02 GB
  2nd save               <killed by OOM killer, no traceback>
  # reproduced twice; both runs died at their second save, neither at the first

after (del state_dict, weight; gc.collect()):
  epoch 3  peak host RSS so far: 8.57GB
  epoch 7  peak host RSS so far: 8.57GB
  epoch 11 peak host RSS so far: 8.57GB
  epoch 15 peak host RSS so far: 8.57GB

Flat across all four saves once the previous dict is released.

Environment Information

• Qwen3-TTS-12Hz-0.6B-Base, qwen-tts at 022e286b98fbec7e1e916cb940cdf532cd9f488e
• Google Colab, Tesla T4 (sm_75), 15360 MiB VRAM
• torch 2.11.0+cu128, transformers 4.57.3
• finetuning/sft_12hz.py, --batch_size 2

Disclosure so you can weigh the report accurately: my run is not the stock configuration. I train with PEFT/LoRA and mixed_precision="no" (the T4 is pre-Ampere, so no bf16), and I call merge_adapter() and filter lora_ keys before saving. But the save path itself is upstream's, the state_dict = {...} rebind at :148 is unmodified, and after merging, the state dict contains the same tensors at the same sizes as a full fine-tune — so defect 2 applies identically to the stock script, and defect 1 is pure control flow and configuration-independent.

Known Issue

• The issue hasn't been already addressed in Documentation, Issues, and Discussions.

I searched the tracker for state_dict, RAM, memory, OOM, killed, disk, atomic, corrupt checkpoint, and copytree before filing. The closest existing reports are #5 and #28 (GPU memory in prepare_data.py, different file and different resource) and #232 / #222 (the speaker_encoder drop in this same block, unrelated to RAM or atomicity). Neither defect above appears to be reported.