Distributed fine-tuning jobs that resume correctly after mid-run pod eviction.
Qwen2.5-0.5B on 2x NVIDIA L4, FSDP2 over NCCL, 500 steps, five mid-run evictions.
step records 500 total, 500 distinct
resumed from [71, 142, 213, 284, 355]
replayed 0 steps
preemptions 5
step 71 save 2088.2 ms grace left 117.8s exit 75
step 142 save 1874.2 ms grace left 118.1s exit 75
step 213 save 2052.5 ms grace left 117.8s exit 75
step 284 save 1894.2 ms grace left 118.0s exit 75
step 355 save 2104.8 ms grace left 117.8s exit 75
every step 1..500 executed across 6 attempts, with resumes
Six attempts, not six pods: this ran under torchrun on a rented two-GPU box
rather than on Kubernetes, because the GPU quota request was refused. Eviction is
the same signal either way, since Kubernetes delivers SIGTERM to PID 1 and that is
what the driver sends. The Kubernetes half is proven separately, on kind, in CI.
| Measure | Value |
|---|---|
| Evictions survived | 5, with zero replayed steps |
| Checkpoint | ~3 GB sharded, 2.05 s median flush against a 120 s grace period |
| Equivalence | bit-identical to an uninterrupted reference, all 291 tensors |
| Nondeterminism floor | 0.000e+00: two identical runs at one seed diverged not at all |
| Throughput | 8,157 tokens/sec; determinism costs about 9% of it |
The reference is the pale line underneath the interrupted one. Five evictions marked, no discontinuity at any.
Zero replayed steps means no work was lost: the handler checkpoints at the step the signal lands on rather than falling back to the last periodic save.
The floor came out at zero, so the interrupted run is bitwise equal to the reference rather than merely close, which is a stronger result than the guarantee below claims for multi-GPU. And bfloat16 resolves to about 0.78% of a value, so the comparison proves the saved parameters agree exactly, not that the arithmetic producing them did.
Full write-up, raw logs, and the reproduction recipe: docs/results.md.
Research compute runs on preemptible capacity because it's 60 to 80% cheaper. That makes checkpoint and resume correctness an economic requirement, not a nicety.
The trap is that a broken resume doesn't look broken. If a job restarts and loses the dataloader position, or the per-rank RNG state, or the optimizer's second moments, the loss curve still goes down. The run just trains on the wrong data in the wrong order. So the failure is invisible without a test that proves equivalence, and that test is the point.
Two tiers.
Tier 1, bit-identical. Single process, torch.use_deterministic_algorithms(True),
CUBLAS_WORKSPACE_CONFIG=:4096:8, fixed seeds. Weights after an
interrupted-and-resumed run match the uninterrupted reference exactly, tensor for
tensor. This runs in CI, on CPU, on every push.
Tier 1 is not limited to one process. A 2-rank FSDP2 run over gloo on CPU is also bit-identical across an interruption, and a test asserts it. Distribution by itself doesn't cost exactness.
Tier 2, equivalent within a measured tolerance. Multi-GPU FSDP over NCCL. NCCL doesn't guarantee a stable reduction order across runs, so exact equality isn't guaranteed there and claiming Tier 1 everywhere would be false.
The obvious framing gets it wrong twice. Tier 2 is not "multi-GPU means
approximate": multi-rank over gloo on CPU keeps Tier 1. And nondeterministic CUDA
kernels aren't a source of tolerance either. use_deterministic_algorithms(True)
raises on an op with no deterministic implementation rather than quietly picking
the fast one, so the run either errors or it's deterministic. That's why the model
pins eager attention: the transformers default is SDPA, whose memory-efficient
backward has no deterministic CUDA kernel.
One caveat on that: use_deterministic_algorithms raises only for ops carrying a
nondeterminism alert, so it's a strong signal and not a proof of exhaustiveness.
kettle measures the tolerance rather than picking one. Running the reference twice at the same seed gives the nondeterminism floor for that hardware and that collective. The assertion is that resuming diverges no more than re-running does. On 2x L4 that floor measured zero, so the interrupted run came out bit-identical.
Losing any one of these breaks equivalence. The last two are the ones that get missed.
| State | Notes |
|---|---|
| Model parameters | Sharded under FSDP |
| Optimizer state | Larger than the model for Adam, and the main checkpoint cost |
| LR scheduler state | Cheap, and easy to forget |
| RNG state, per rank | Python, NumPy, torch CPU, and torch CUDA. Per rank, not global |
| Dataloader position | Epoch, shuffle seed, and exact sample index |
Gradient accumulation position is deliberately not on that list. Checkpoints are
only ever taken at optimizer step boundaries, after optimizer.step() and
zero_grad(), so a preemption arriving mid-accumulation discards those gradients
rather than persisting them. That costs at most one step of work per eviction and
removes an entire class of partial state from the design.
Two of those five are subtler than they look. So is one property of the loop that isn't state at all.
The dataloader. StatefulDataLoader captures position within the current
iterator, but the per-epoch reshuffle is driven by a generator that isn't in the
state dict. So resume inside an epoch is correct and resume across an epoch
boundary silently isn't. kettle deletes the hidden state rather than chasing it:
each epoch builds its loader from a generator seeded by (seed, epoch), so the
permutation is a pure function of the config and epoch becomes an explicit
checkpointed integer.
Per-rank state gets deduplicated. torch.distributed.checkpoint deduplicates
on the state-dict key alone and never compares values, so two ranks writing
different things under one key collapse to a single copy, and the survivor is
whichever rank's write plan is smallest. That's intended: DCP treats a shared key
as replicated state. Nothing checks the assumption. Under a plain rng key a
2-rank checkpoint held a single RNG state and a single dataloader position, so
both ranks restored the same ones and their streams correlated. Worth about
4e-4 of weight divergence: enough to fail a bit-identical assertion, nowhere
near enough to see in a loss curve. Per-rank entries are keyed rng.rank0,
loader.rank0 and so on, which makes them un-collidable.
The stop decision has to be collective. A rank that sees SIGTERM and decides on its own to checkpoint enters a collective its peers never reach, and blocks until the grace period kills it. The job doesn't checkpoint, it hangs. So the preemption flag is all-reduced inside the accumulation loop, and one rank seeing the signal stops all of them at the same point by agreement rather than by luck.
- Run a reference: N steps, no interruption, seed S.
- Run a trial with the same config and the same seed. SIGTERM it at step K.
- Restart. It resolves the latest complete checkpoint and continues to N.
- Assert the trial's weights at step N match the reference's at step N.
- Assert the batch consumed at step K+1 is identical in both runs.
Every restored piece of state also gets a negative control. Each one disables the restoration of exactly one thing and requires the equivalence check to notice:
test_resumed_run_is_bit_identical_to_the_reference PASSED
test_resume_actually_happened PASSED
test_control_skipping_rng_restore_breaks_equivalence PASSED
test_control_skipping_dataloader_restore_breaks_equivalence PASSED
test_control_skipping_optimizer_restore_breaks_equivalence PASSED
test_control_skipping_scheduler_restore_breaks_equivalence PASSED
test_resume_falls_back_past_a_torn_write PASSED
The controls matter more than the positive test. A resume that quietly restores nothing still produces a falling loss curve, and it would sail through the positive test if the thing being restored never mattered in the first place.
That isn't hypothetical. At dropout 0 the training loop draws from no global RNG stream at all, so a resume restoring no RNG state would pass perfectly and the entire dataloader and RNG effort would be untested. Dropout is a hashed config field, and a test asserts the forward pass actually advances the RNG.
Runs in CI on every push: kind, MinIO, Kueue, then a running pod deleted mid-run.
==> at step 376, evicting
==> resumed in kettle-train-8sckg
==> at step 751, evicting
==> resumed in kettle-train-jvk8h
==> verifying
step records 1500 total, 1500 distinct
resumed from [384, 757]
replayed 0 steps
preemptions 2
step 384 save 105.6 ms grace left 29.9s exit 75
step 757 save 107.1 ms grace left 29.9s exit 75
every step 1..1500 executed across 3 attempts, with resumes
Zero replayed steps again, this time on a platform that did the evicting
itself. Without the SIGTERM handler each eviction would silently cost up to
checkpoint_every steps and the run would still finish looking correct.
The same driver, scripts/chaos.py, produced the GPU result at the top. It takes
either a Kubernetes overlay or a config, so the free test and the paid one run
identical code and only the thing pulling the trigger differs.
Eviction points are computed, never improvised: evenly spaced by default,
jittered within the same bands under --seed, and never past 80% of the run. An
eviction with no work behind it, or one landing near the finish line, produces a
run that passes while proving nothing. Those rules are tests, not comments.
A Job that counts evictions against backoffLimit eventually gives up on a run
that was working perfectly. podFailurePolicy is how Kubernetes gets told the
difference:
DisruptionTargetcovers eviction by the platform: node drain, preemption by a higherPriorityClass, spot reclaim.- Exit 75 is kettle's own "I checkpointed and stepped aside".
Both are Ignore, so neither counts. Anything else does, so a genuine crash loop
still terminates instead of burning quota forever.
podFailurePolicy requires restartPolicy: Never, which is why the Job doesn't
use OnFailure. That's the better fit anyway: each attempt gets a fresh pod,
which is what a real preemption looks like, and per-attempt logs stay separable.
A Kubernetes Job reads exactly one bit from a container: did it exit zero. So
exit 0 has to mean the run reached total_steps, and nothing else. A container
that checkpoints on SIGTERM and then exits cleanly would mark an unfinished Job
complete and never come back.
- Exit 0 only at
total_steps. - Exit 75 after a preemption checkpoint, so the Job restarts. That's
EX_TEMPFAILfromsysexits.h, a temporary failure where retry is appropriate. Rejected 143, the conventional128 + SIGTERM, because that means a process was killed without handling the signal, which is the opposite of what happened. backoffLimithas to exceed the number of preemptions the run is expected to survive.
The container has to actually report that code, which is less automatic than it sounds: torchrun collapses any non-zero worker exit into 1. So the entrypoint keeps the shell as PID 1, forwards SIGTERM, and reports 75 itself. CI asserts all three paths, because a fix that made everything look like a preemption would be worse than the bug: 75 on preemption, 0 on clean completion, 1 on a crash.
Deleting a pod deletes its logs, and the preemption record is written during the grace period, after the delete. So the harness follows every pod as it appears and streams into one file. Reading logs at the end would lose exactly the evidence the test is about.
terminationGracePeriodSeconds is a hard ceiling. When it expires the container
gets SIGKILL, and if the checkpoint hasn't landed the work is lost no matter how
correct the code is. So the number that matters isn't the save duration on its
own, it's the save duration against the grace period.
kettle emits both on every preemption, in the same JSONL stream as the step metrics, so every number here came out of a run that was actually evicted. On Qwen2.5-0.5B across 2x L4 that is a ~3 GB sharded flush in 2.05 s against 120 s, about 58x headroom. On the 107K-parameter CI model it is 15 ms against 30 s, which is true and tells you nothing: the ratio only becomes a design constraint at a real model size.
The budget has three numbers, not two. terminationGracePeriodSeconds is what
Kubernetes allows and grace_seconds is what the process budgets against, but
the container runs under torchrun, which has its own --shutdown-timeout
defaulting to 30 seconds. On SIGTERM it signals its workers and SIGKILLs them
when that expires, regardless of what the platform allows. A pod granted 120 seconds would still lose its workers at 30, mid-write,
on every eviction: the run stays correct because an uncommitted checkpoint is
ignored, and quietly loses up to checkpoint_every steps each time. The grace
period travels as KETTLE_GRACE from the manifest into torchrun, and the test is
a three-way assertion.
Two things fall out of this:
- The save is attempted even when the remaining grace looks too short. A checkpoint SIGKILLed partway leaves an uncommitted directory, which
latest_completealready ignores, so the downside of trying is losing exactly the work that would have been lost anyway. A direct payoff from writing_COMMITlast. - A preemption landing on an already-checkpointed step exits immediately. Nothing is spent re-saving state that's already durable.
The correctness work runs on CPU, in CI, for free, and that constraint is what
makes it affordable. It also has a blind spot. A tensor on the wrong
device costs nothing when there's one device. A collective that allocates on CPU
works under gloo. use_deterministic_algorithms(True) raises on CUDA where it
silently permits on CPU.
An audit before the first GPU run found six defects of exactly that shape:
| What | How it would have shown up |
|---|---|
| The CUDA image resolved to a torch two minor versions below the lockfile | Every pod exits on an unrecognized argument before training a step |
| Model and batches never moved to the device | Device mismatch at the first embedding lookup, or on one GPU, silent CPU training |
any_rank and mean all-reduced CPU tensors |
NCCL has no CPU backend, so it raises on the first micro-batch |
| ClusterQueue allowed 16Gi while the Job requested 32Gi | Kueue never admits, the Job stays suspended, no pod is created |
torchrun's --shutdown-timeout defaults to 30s |
Workers SIGKILLed mid-checkpoint on every eviction, losing work silently |
| transformers defaults to SDPA, whose backward has no deterministic CUDA kernel | Raises at the first backward pass |
Most would have been caught by one five-step run on a rented L4, for about twenty
cents. tests/test_gpu.py is that run: device placement, NCCL collectives on the
real backend, deterministic kernels against real CUDA, CUDA RNG through a
checkpoint, 2-rank FSDP2 save and resume, and a 1 to 2 reshard. It asserts nothing
about loss. It asserts that the machinery touches the hardware it claims to, and
it runs before terraform apply.
The missing rung was between free CPU tests and a GPU cluster, and it cost about a dollar to add.
| Decision | Choice | Why, and what was rejected |
|---|---|---|
| Framework | PyTorch FSDP2 (fully_shard) |
First-party, fewer moving parts. Rejected DeepSpeed: more config surface, and ZeRO-3 is comparable at this scale. |
| Checkpointing | torch.distributed.checkpoint |
Built for sharded save and load, and for resharding onto a different GPU count. Rejected gathering state_dict() to rank 0: it OOMs and serializes. |
| Dataloader | torchdata StatefulDataLoader |
Real resumable iteration state. Rejected re-seeding and fast-forwarding: it's O(n), and it's wrong with multiple workers. |
| Run identity | SHA-256 over the computation-affecting config fields | A run can only load a checkpoint written by a config that computes the same thing. Rejected hashing the whole config: moving buckets would change the run ID and orphan every checkpoint already written. |
| Restart | restartPolicy: Never plus a podFailurePolicy |
A fresh pod per attempt, and evictions don't count against backoffLimit. Rejected OnFailure, which counts every preemption as a failure and eventually abandons a healthy run. |
| Checkpoint completeness | A _COMMIT marker written last |
Object stores have no atomic directory rename, so a listing can show a half-written checkpoint. Rejected trusting the listing. |
| Queueing | Kueue, quota and admission only | GPU quota enforcement for a few lines of YAML. Rejected writing a scheduler, and rejected an operator or CRD. |
| Storage | fsspec URIs | file://, s3://, and gs:// become one code path. Rejected a storage abstraction layer with one implementation. |
| Parallelism | FSDP2, though DDP would fit | Chosen for the sharded checkpoint and its resharding behavior, not for memory. See below. |
| Cloud | Terraform for GKE and a spot GPU node pool | The platform-specific surface stays in infra/ and a Kustomize overlay. |
The easy assumption is wrong. FSDP shards parameters, gradients, and optimizer state across ranks; it does not split the math of a forward pass. Each rank runs the whole model on its own microbatch, exactly like DDP. The parallelism is plain data parallelism, and sharding is an orthogonal trick for fitting a model that wouldn't otherwise fit.
Neither model here needs it. The CI model's entire training footprint is under 2 MB, and Qwen2.5-0.5B in bf16 is about 1 GB against 23 GB of L4.
FSDP2 is here because the interesting artifact is the sharded checkpoint and
its resharding behavior: a DTensor state dict, one shard per rank, portable
across a topology change.
The decision rule: when the model fits replicated, DDP is faster
and simpler, since it does one all_reduce per step against FSDP's all_gather
per layer in both directions plus a reduce_scatter. You reach for FSDP when
replication doesn't fit, not for speed.
file://, s3://, and gs:// are the same code path, because fsspec already
abstracts them. There's no StorageBackend base class with one implementation.
What doesn't abstract away is that object stores have no directories. A step
isn't a folder, it's a shared key prefix. So nothing asks whether a directory
exists, and completeness is answered by globbing for the marker objects
themselves, which is one request instead of one listing plus a probe per step. A
test asserts that resolving the newest checkpoint costs exactly one glob.
The object-store tests run against a live S3 API on every push, via moto
in-process, so they need no Docker daemon and no service container. Setting
KETTLE_S3_ENDPOINT runs the identical tests against MinIO instead.
These are firm:
- Training from scratch, novel architectures, or competitive model quality
- Beating anyone's throughput numbers
- Multi-node training
- A Kubernetes operator, CRD, or admission webhook
- A web UI, auth, or multi-tenancy
- Hyperparameter search, distillation, RLHF, or quantization
- Supporting more than one model family
make install # uv sync
make test-fast # everything except the torchrun tests
make test-dist # 2-rank FSDP2 over gloo, still CPU-only
make lint
make image # CPU image; TORCH_GROUP=cu129 builds CUDA + linux/amd64
make e2e-local # kind + MinIO + Kueue, evict a pod mid-run, verify the resume
make test-gpu # the pre-flight suite. Needs two GPUs.Submitting a run:
$ kettle --overlay deploy/overlays/kind submit
run_id kind-e2e-ae5bc7e54e1a
fingerprint ae5bc7e54e1ad4583adfcf9908f4325f34eafdc4d7cb11bdd632e9e5fc9edd93
checkpoints s3://kettle-checkpoints/runs/kind-e2e-ae5bc7e54e1a
steps 1500
kettle submit renders the overlay and applies it. It exists rather than a bare
kubectl apply -k because of the run ID: submitting the same experiment twice
lands on the same checkpoint prefix, so the second submission resumes the first.
kettle status reads both the Job and the checkpoint store, since those can
disagree and the gap between them is the interesting number after an eviction.
kettle evict deletes the running pod the way a preemption would.
Training emits one JSONL record per optimizer step on stdout, carrying the run
ID, loss, learning rate, and throughput. Structured output, because it comes back
through kubectl logs.
On GCP, make gke-up runs terraform apply, installs Kueue, writes the cluster's
values into a gitignored copy of the overlay, and pushes the CUDA image. It needs
gcloud authenticated and GPU quota in the target region. make gke-destroy
removes all of it. What a cloud has to provide, and how a second provider would
slot in, is infra/README.md.
Reproducing the GPU run without a cluster:
docs/results.md.
- Elastic scaling:
torchrunwith a min and max world size, so a run shrinks onto surviving capacity instead of stopping. This needs the existing resharding path. - Spot interruption notices: GCP publishes preemption on the metadata endpoint roughly 30 seconds ahead of shutdown. Reacting to that beats reacting to SIGTERM, since the checkpoint starts before the grace period clock does.
- Multi-node: where NCCL failure modes get genuinely interesting, and where a partially scheduled job needs real gang scheduling rather than the single-pod version of it.
- Asynchronous checkpointing: the grace period budget caps how large a checkpoint can get. Sharded async save pushes that ceiling up, and quantifying where it lands is the more useful version of this project.
- Bake the model weights into the image:
from_pretrainedcaches into the container's writable layer, which dies with the pod, so a run surviving five evictions downloads the model six times and depends on the hub staying reachable.