Fix RamTorch prefetch-order save race

[hjinnkim]

Summary

SaveHookManager.save_model_hook writes the shared RamTorch prefetch-order file before
is_main_process is resolved:

# save_hooks.py:1094-1098 on main
StateTracker.save_training_state(os.path.join(output_dir, self.training_state_path))
StateTracker.save_ramtorch_prefetch_orders(output_dir)          # <- every participating rank

distributed_type = DistributedType.NO
is_main_process = True

save_training_state writes a rank-specific filename (training_state-rank{rank}.json,
save_hooks.py:330-334) and is correct as-is. The prefetch orders are one shared file per
checkpoint directory, written by StateTracker.save_ramtorch_prefetch_orders
(state_tracker.py:402-416):

temp_path = path.with_suffix(f"{path.suffix}.tmp")     # same name on every rank
with temp_path.open("w") as handle:
    fcntl.flock(handle, fcntl.LOCK_EX)
    ...
    fcntl.flock(handle, fcntl.LOCK_UN)
temp_path.replace(path)                                # outside the lock

There is an flock here, which is why this looks safe at a glance — but it does not make the
operation atomic. The temp name contains no rank, so every rank opens the same path with "w"
and truncates whatever the others wrote; the lock is released before the handle closes; and the
replace() is outside the lock entirely. Whichever process reaches replace() after the first has
already renamed the temp file away finds nothing there and raises FileNotFoundError.

This PR makes each node's local main process the owner of the write
(is_local_main_process), so the file also lands on node-local storage in multi-node runs, and
gives the writer a process-unique temp filename (.tmp.{pid} + atomic rename) so concurrent
node-mains on a shared filesystem cannot race on the rename. (The first commit gated the write on
the global main process; the second widened it to one writer per node after the multi-node
discussion below.)

Which configurations are affected

Plain DDP is not affected. The only entry point to the registered save-state pre-hook is
accelerator.save_state at trainer.py:5982, and its callers gate on
is_main_process or use_deepspeed_optimizer or fsdp_enable (trainer.py:5251) and
is_main_process or use_deepspeed_optimizer for the rolling path (trainer.py:6824). Under DDP
only rank 0 ever reaches save_model_hook, so there is nothing to race. The defect is reachable
with use_deepspeed_optimizer=true or fsdp_enable=true, where every rank saves.

Who fails

Not just the non-main ranks. Measured over 60 barrier-synchronised iterations, rank 0 raised on
44 of 60 — statistically indistinguishable from ranks 1-7 (41-50 of 60). The loser is whichever
process calls replace() second, regardless of rank.

What the failure costs

Measured: 8 concurrent processes calling StateTracker.save_ramtorch_prefetch_orders against a
shared directory produced 1206 FileNotFoundError out of 1600 calls, all at
state_tracker.py:416.

That the whole job goes down is reasoned, not measured: accelerate invokes save-state pre-hooks
with no try/except, and the accelerator.save_state call at trainer.py:5982 sits outside any
enclosing try in checkpoint_state_save (trainer.py:5933). The function contains two try
blocks and both are closed before that line — :5958-5966 catches Exception around debug-logging
of the registered model list, and :5978-5981 catches FileNotFoundError around removal of the
checkpoint guard file, inside the is_main_process branch. Neither wraps the save, so the
exception propagates out of the training loop on whichever rank raised it. The observed
non-determinism is consistent with an intermittent hard failure at checkpoint time.

Changes

• Keep the existing per-rank StateTracker.save_training_state() call exactly where it is.
• Resolve is_local_main_process before writing the shared RamTorch prefetch-order file, and call
  StateTracker.save_ramtorch_prefetch_orders() only from each node's local main process — one
  writer per node, so multi-node runs with node-local output_dir keep a copy on every node.
• Give the writer a process-unique temp filename (.tmp.{pid}), making the final replace() an
  atomic last-writer-wins rename even when two node-mains share a filesystem.
• Regression coverage: local-main gating (including the is_main_process=False /
  is_local_main_process=True case that models a non-zero node), per-rank training-state writes,
  and a two-process barrier-synchronised concurrency test against the real writer.

Notes

• This changes ownership of one shared write. Rank-specific training-state filenames and all
  restore paths are unchanged, as is the checkpoint format.
• Resume still reads the file on every rank; only the writer is narrowed.
• On a shared filesystem the node-mains' copies are equivalent (each process learns the same
  transition statistics), so last-writer-wins does not change behaviour.
• No GPU or live distributed run was performed for this branch. The race was reproduced with
  concurrent CPU processes against the real StateTracker writer; the multi-node layouts were
  simulated the same way (see the review thread for the full script and measurements).

Verification

CPU-only container, current main (2d6df1245).

Commit 1 (25cd2460, single-writer gate) — RED is main plus this branch's test file with
save_hooks.py restored from main:

tests.test_save_hooks       RED   Ran 1    FAILED (failures=1)
                            GREEN Ran 1    OK

related modules (133)       RED   Ran 133  FAILED (failures=1, skipped=7)
                            GREEN Ran 133  OK (skipped=7)

46 modules (515 tests)      RED   Ran 515  FAILED (failures=1, errors=1, skipped=18)
                            GREEN Ran 515  FAILED (errors=1, skipped=18)
                            GREEN-only failures: NONE

Commit 2 (2891e79d, per-node writer + unique temp) — RED is the previous head 25cd2460
plus the new test file, i.e. the tests are pinned to this commit's production delta:

tests.test_save_hooks       RED   Ran 5    FAILED (failures=2)
                            GREEN Ran 5    OK

46 modules (515 tests)      GREEN Ran 515  FAILED (errors=1, skipped=18)   # identical to baseline

The two RED failures are exactly the tests that encode this commit:
test_ramtorch_prefetch_orders_ignore_global_main_flag (a non-zero node's local main must write)
and test_concurrent_writers_do_not_race (two real processes, 30 barrier-synchronised rounds —
red with the shared temp name, green with the process-unique one). Driving the production writer
through the 2-node simulation from the review thread: 0 errors, the file present on every node
(node-local layout) and valid JSON with no leftover temp files (shared layout).

The one error present on both trees is environmental — ModuleNotFoundError: No module named 'torchaudio' in tests.test_acestep_lora_targets — and is unrelated to this change.

[bghira]

but is_main_process wouldn't happen to save it on each node in multi-node training runs. can you verify that?

[Copilot]

Pull request overview

Fixes a distributed checkpoint-time race by ensuring the shared RamTorch prefetch-order file is written only by the main process, after is_main_process is determined inside SaveHookManager.save_model_hook.

Changes:

• Move StateTracker.save_ramtorch_prefetch_orders(output_dir) to after is_main_process resolution and guard it with if is_main_process:.
• Add a unit test covering is_main_process=True/False behavior for the RamTorch prefetch-order save call.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
simpletuner/helpers/training/save_hooks.py	Guards the shared RamTorch prefetch-order write so only the main process performs it, eliminating the multi-rank temp-file rename race.
tests/test_save_hooks.py	Adds a regression test asserting the prefetch-order file write occurs only on the main process.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[hjinnkim]

Confirmed — with is_main_process the file is written only on node 0's disk. Measured: after the simulation below, node0: True, node1: False.

How this was measured. No multi-node cluster was involved: 8 Python processes stand in for 2 nodes x 4 ranks. "Node-local storage" is modelled as one output directory per node (ranks 0-3 write to node0/, ranks 4-7 to node1/); a shared filesystem as a single directory for all 8. Each save cycle is barrier-synchronised so all permitted writers hit StateTracker.save_ramtorch_prefetch_orders (the real writer, imported from this repo) at the same moment, 50 cycles per case. The writer-selection policies compared are all (current main), global_main (this PR), and local_main (one writer per node). Script is in the collapsed section below.

Results, and what they mean for the question:

1. This commit is focused on the critical failure, and it does resolve it. With the all-ranks writer on main, the write itself crashes before any node gets a usable copy: 120/200 and 142/200 save calls per node raised FileNotFoundError (all ranks share one temp filename and replace() races). That takes the job down at checkpoint time. With the single-writer commit as it stands: 0 errors in every layout tested.
2. Resume on a node without the file does not fail. Every rank calls StateTracker.load_ramtorch_prefetch_orders on restore; a missing file hits the path.exists() guard (state_tracker.py:378) and returns an empty learned state with no exception — verified by calling the loader against the file-less node1/ directory after the run. Prefetch falls back to the static traversal order and re-learns (3 observations / 0.80 confidence per transition) — a short warmup on those nodes, nothing else.

If each node should keep its own copy: gating the write on is_local_main_process with a process-unique temp name (.tmp.{process_index} + atomic rename) measured 0 errors in both node-local and shared-filesystem layouts, with the file present on every node. I can push that variant to this branch with the regression test extended to cover both layouts.

Reproduction script and raw output

Run with PYTHONPATH pointing at the repo checkout (CPU only, no GPU needed):

"""Multi-node simulation for the is_main_process question on PR #2988."""
import json
import multiprocessing as mp
import sys
import tempfile
from pathlib import Path

N_NODES = 2
LOCAL_WORLD = 4
ITERS = 50


def _unique_tmp_save(state_tracker, directory, rank):
    """The proposed writer: process-unique temp name + atomic rename."""
    path = state_tracker._ramtorch_prefetch_order_path(directory)
    data = state_tracker._normalise_ramtorch_prefetch_orders(state_tracker.ramtorch_prefetch_orders)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(f"{path.suffix}.tmp.{rank}")
    with tmp.open("w") as handle:
        json.dump(data, handle)
    tmp.replace(path)


def worker(rank, node_dirs, gate, unique_tmp, barrier, errq):
    from simpletuner.helpers.training.state_tracker import StateTracker

    node = rank // LOCAL_WORLD
    local_rank = rank % LOCAL_WORLD
    outdir = node_dirs[node]
    allowed = {
        "all": True,                    # main today: every rank writes
        "global_main": rank == 0,       # this PR: is_main_process
        "local_main": local_rank == 0,  # proposed: is_local_main_process
    }[gate]

    errors = 0
    for _ in range(ITERS):
        barrier.wait()
        if not allowed:
            continue
        try:
            if unique_tmp:
                _unique_tmp_save(StateTracker, outdir, rank)
            else:
                StateTracker.save_ramtorch_prefetch_orders(outdir)
        except FileNotFoundError:
            errors += 1
    errq.put((rank, node, errors))


def run_case(name, gate, unique_tmp, shared_fs):
    base = Path(tempfile.mkdtemp(prefix=f"probe-{name}-"))
    node_dirs = [base / "shared"] * N_NODES if shared_fs else [base / f"node{n}" for n in range(N_NODES)]
    for d in node_dirs:
        d.mkdir(parents=True, exist_ok=True)

    world = N_NODES * LOCAL_WORLD
    barrier = mp.Barrier(world)
    errq = mp.Queue()
    procs = [mp.Process(target=worker, args=(r, node_dirs, gate, unique_tmp, barrier, errq)) for r in range(world)]
    for p in procs:
        p.start()
    for p in procs:
        p.join()

    per_node_errors = {}
    for _ in range(world):
        rank, node, errors = errq.get()
        per_node_errors[node] = per_node_errors.get(node, 0) + errors
    files = {
        ("shared" if shared_fs else f"node{n}"): (node_dirs[n] / "ramtorch_prefetch_orders.json").exists()
        for n in range(N_NODES)
    }
    print(f"=== {name}  gate={gate} unique_tmp={unique_tmp} {'shared-FS' if shared_fs else 'node-local'} ===")
    print(f"  FileNotFoundError per node : {per_node_errors}")
    print(f"  file present after run     : {files}")
    return node_dirs


def main():
    mp.set_start_method("spawn")
    run_case("case1-main-today", gate="all", unique_tmp=False, shared_fs=False)
    node_dirs = run_case("case2-this-pr", gate="global_main", unique_tmp=False, shared_fs=False)

    from simpletuner.helpers.training.state_tracker import StateTracker
    StateTracker.reset_ramtorch_prefetch_orders()
    loaded = StateTracker.load_ramtorch_prefetch_orders(node_dirs[1])
    print(f"  resume on node1 (no file)  : load returned {loaded!r} (no exception)")

    run_case("case3-local-main", gate="local_main", unique_tmp=True, shared_fs=False)
    run_case("case4-local-main-sharedfs", gate="local_main", unique_tmp=True, shared_fs=True)


if __name__ == "__main__":
    main()

Output (error counts vary run to run; this is one run verbatim):

=== case1-main-today  gate=all unique_tmp=False node-local ===
  FileNotFoundError per node : {0: 120, 1: 142}
  file present after run     : {'node0': True, 'node1': True}
=== case2-this-pr  gate=global_main unique_tmp=False node-local ===
  FileNotFoundError per node : {1: 0, 0: 0}
  file present after run     : {'node0': True, 'node1': False}
  resume on node1 (no file)  : load returned {'version': 1, 'components': {}} (no exception)
=== case3-local-main  gate=local_main unique_tmp=True node-local ===
  FileNotFoundError per node : {0: 0, 1: 0}
  file present after run     : {'node0': True, 'node1': True}
=== case4-local-main-sharedfs  gate=local_main unique_tmp=True shared-FS ===
  FileNotFoundError per node : {1: 0, 0: 0}
  file present after run     : {'shared': True}

Note case1 shows the file "present" only because the final cycle's survivor completed a rename — 60%+ of the save calls before it raised, any one of which ends a real training job.

[bghira]

yeah let's be sure multi-node training isn't broken by any PRs submitted, thank you

[hjinnkim]

Pushed 2891e79d: the write is now gated on is_local_main_process (one writer per node, so the file lands on node-local storage in multi-node runs) and the writer uses a process-unique temp filename (.tmp.{pid} + atomic rename), which also removes the remaining collision between node-mains on a shared filesystem. Driving the production writer through the 2-node simulation above: 0 errors in both layouts, file present on every node, no leftover temp files.

Regression tests added: a non-zero node's local main (is_main_process=False, is_local_main_process=True) must write, and a two-process barrier-synchronised concurrency test against the real writer — both red on the previous head, green now. 46-module suite unchanged from baseline. PR description updated to match.